Longest subarray with given conditions

Given 2 integers S and K and an array arr[] of N integers, find the longest subarray in arr[] such that the subarray can be divided into at most K contiguous segments and the sum of elements of each segment can be at most S.

Example:

Input: n = 5, k = 2, s = 5, arr = { 1, 3, 2, 1, 5 }
Output: 4
Explanation: One of the segments with length 4 is [1,3,2,1]. It can be divided into two segments [1,3] and [2,1] whose sums are at most S=5.

Input: n = 5, k = 3, s = 5, arr = {5, 1 ,5 ,1 ,1}
Output: 4
Explanation: One of the segments with length 4 is [1,5,1,1]. It can be divided into three segments [1], [5] and [1,1] whose sums are at most S=5.

Approach:

The idea is to uses a two-pointer approach to find the right boundary of each segment, precomputes right boundaries using binary lifting, and make use of binary lifting to determine the maximum number of elements in a good segment for each starting index. The result is the maximum value encountered during these calculations.

• Initialize Variables:
  • Set up variables for the current sum (currSum) and the maximum number of elements in a good segment (maxGoodSegment).
• Segment Boundary Calculation:
  • Use a two-pointer approach to find the right boundary of each segment:
  • Iterate over possible left boundaries.
  • Extend the right boundary until the sum exceeds s.
  • Store the right boundary in the 2D array up.
• Precompute Right Boundaries:
  • Use binary lifting to precompute the right boundaries for each power of 2.
• Binary Lifting for Maximum Segments:
  • Use binary lifting to find the maximum number of elements in a good segment for each starting index (i):
  • Iterate over the binary representation of k.
  • Update the current index based on the binary lifting technique.
  • Track the maximum segment size encountered during these iterations.
• Output Result:
  • The maximum number of elements in a good segment (maxGoodSegment) represents the solution for the given input

Below is the implementation of the above approach:

C++

#include <bits/stdc++.h>
using namespace std;

const int MAX_N = 1e5;

int up[MAX_N + 5][20];

int solve(int n, int k, int s, vector<int>& arr)
{

    vector<int> nums(n, 0);
    int currSum = 0;
    int maxGoodSegment = 0;

    // Calculate the maximum number of elements in a good
    // segment
    for (int i = 0, j = 0; i < n; i++) {
        // Extend the current segment until the sum exceeds
        // 's'
        while (j < n && currSum + arr[j] <= s) {
            currSum += arr[j];
            j++;
        }

        // Store the right boundary of the current segment
        up[i][0] = j;
        currSum -= arr[i]; // Move the left boundary of the
                           // segment
    }

    up[n][0] = n;

    // Precompute the right boundaries for each power of 2
    for (int j = 1; j < 20; j++) {
        for (int i = 0; i <= n; i += 1) {
            up[i][j] = up[up[i][j - 1]][j - 1];
        }
    }

    // Calculate the maximum number of elements in a good
    // segment using binary lifting
    for (int i = 0; i < n; i++) {
        int currIndex = i;
        for (int j = 19; j >= 0; j--) {
            if (k & (1 << j)) {
                currIndex = up[currIndex][j];
            }
        }
        maxGoodSegment = max(maxGoodSegment, currIndex - i);
    }

    // Output the result for the current test case
    return maxGoodSegment;
}

int main()
{

    int n = 5, k = 2, s = 5;

    vector<int> arr = { 1, 3, 2, 1, 5 };

    cout << solve(n, k, s, arr);

    return 0;
}

Java

import java.util.Arrays;
import java.util.Vector;

public class Main {
    static final int MAX_N = 100005;
    static int[][] up = new int[MAX_N + 5][20];

    static int solve(int n, int k, int s, Vector<Integer> arr) {
        // Vector to store nums array, initialized with 0
        Vector<Integer> nums = new Vector<>(n);
        for (int i = 0; i < n; i++) {
            nums.add(0);
        }

        int currSum = 0;
        int maxGoodSegment = 0;

        // Calculate the right boundaries of segments
        for (int i = 0, j = 0; i < n; i++) {
            while (j < n && currSum + arr.get(j) <= s) {
                currSum += arr.get(j);
                j++;
            }
            up[i][0] = j; // Store the right boundary of the current segment
            currSum -= arr.get(i); // Move the left boundary of the segment
        }

        up[n][0] = n; // Set the right boundary of the last element

        // Precompute the right boundaries for each power of 2 using binary lifting
        for (int j = 1; j < 20; j++) {
            for (int i = 0; i <= n; i += 1) {
                up[i][j] = up[up[i][j - 1]][j - 1];
            }
        }

        // Calculate the maximum number of elements in a good segment using binary lifting
        for (int i = 0; i < n; i++) {
            int currIndex = i;
            for (int j = 19; j >= 0; j--) {
                if ((k & (1 << j)) != 0) {
                    currIndex = up[currIndex][j];
                }
            }
            maxGoodSegment = Math.max(maxGoodSegment, currIndex - i);
        }

        return maxGoodSegment;
    }

    public static void main(String[] args) {
        // Example input
        int n = 5, k = 2, s = 5;
        Vector<Integer> arr = new Vector<>(Arrays.asList(1, 3, 2, 1, 5));

        // Output the result for the current test case
        System.out.println(solve(n, k, s, arr));
    }
}

// This code is contributed by akshitaguprzj3

Python3

MAX_N = 10**5

up = [[0] * 20 for _ in range(MAX_N + 5)]

def solve(n, k, s, arr):
    nums = [0] * n
    curr_sum = 0
    max_good_segment = 0

    # Calculate the maximum number of elements in a good segment
    for i in range(n):
        j = i
        # Extend the current segment until the sum exceeds 's'
        while j < n and curr_sum + arr[j] <= s:
            curr_sum += arr[j]
            j += 1

        # Store the right boundary of the current segment
        up[i][0] = j
        curr_sum -= arr[i]  # Move the left boundary of the segment

    up[n][0] = n

    # Precompute the right boundaries for each power of 2
    for j in range(1, 20):
        for i in range(n + 1):
            up[i][j] = up[up[i][j - 1]][j - 1]

    # Calculate the maximum number of elements in a good segment using binary lifting
    for i in range(n):
        curr_index = i
        for j in range(19, -1, -1):
            if k & (1 << j):
                curr_index = up[curr_index][j]
        max_good_segment = max(max_good_segment, curr_index - i)

    # Output the result for the current test case
    return max_good_segment

if __name__ == "__main__":
    n, k, s = 5, 2, 5
    arr = [1, 3, 2, 1, 5]

    print(solve(n, k, s, arr))

C#

using System;
using System.Collections.Generic;

public class Solution
{
    const int MAX_N = 100005;

    static int[,] up = new int[MAX_N + 5, 20];

    static int Solve(int n, int k, int s, List<int> arr)
    {
        int currSum = 0;
        int maxGoodSegment = 0;

        // Calculate the maximum number of elements in a good
        // segment
        for (int i = 0, j = 0; i < n; i++)
        {
            // Extend the current segment until the sum exceeds
            // 's'
            while (j < n && currSum + arr[j] <= s)
            {
                currSum += arr[j];
                j++;
            }

            // Store the right boundary of the current segment
            up[i, 0] = j;
            currSum -= arr[i]; // Move the left boundary of the
                               // segment
        }

        up[n, 0] = n;

        // Precompute the right boundaries for each power of 2
        for (int j = 1; j < 20; j++)
        {
            for (int i = 0; i <= n; i += 1)
            {
                up[i, j] = up[up[i, j - 1], j - 1];
            }
        }

        // Calculate the maximum number of elements in a good
        // segment using binary lifting
        for (int i = 0; i < n; i++)
        {
            int currIndex = i;
            for (int j = 19; j >= 0; j--)
            {
                if ((k & (1 << j)) != 0)
                {
                    currIndex = up[currIndex, j];
                }
            }
            maxGoodSegment = Math.Max(maxGoodSegment, currIndex - i);
        }

        // Output the result for the current test case
        return maxGoodSegment;
    }

    public static void Main(string[] args)
    {
        int n = 5, k = 2, s = 5;

        List<int> arr = new List<int> { 1, 3, 2, 1, 5 };

        Console.WriteLine(Solve(n, k, s, arr));
    }
}

Javascript

// Javascript code for the above approach

const MAX_N = 100005;
const up = Array.from({
    length: MAX_N + 5
}, () => Array.from({
    length: 20
}, () => 0))

function solve(n, k, s, arr) {
    // Array to store nums array, initialized with 0
    const nums = Array(n).fill(0);

    let currSum = 0;
    let maxGoodSegment = 0;

    // Calculate the right boundaries of segments
    for (let i = 0, j = 0; i < n; i++) {
        while (j < n && currSum + arr[j] <= s) {
            currSum += arr[j];
            j++;
        }
        up[i][0] = j; // Store the right boundary of the current segment
        currSum -= arr[i]; // Move the left boundary of the segment
    }

    up[n][0] = n; // Set the right boundary of the last element

    // Precompute the right boundaries for each power of 2 using binary lifting
    for (let j = 1; j < 20; j++) {
        for (let i = 0; i <= n; i += 1) {
            up[i][j] = up[up[i][j - 1]][j - 1];
        }
    }

    // Calculate the maximum number of elements in a good segment using binary lifting
    for (let i = 0; i < n; i++) {
        let currIndex = i;
        for (let j = 19; j >= 0; j--) {
            if ((k & (1 << j)) != 0) {
                currIndex = up[currIndex][j];
            }
        }
        maxGoodSegment = Math.max(maxGoodSegment, currIndex - i);
    }

    return maxGoodSegment;
}

// Example input
const n = 5,
    k = 2,
    s = 5;
const arr = [1, 3, 2, 1, 5];

// Output the result for the current test case
console.log(solve(n, k, s, arr));

// This code is contributed by ragul21

Output

4

Time complexity: O(n + n * log(k)), where n is the size of the input array and k is the given parameter.
Auxiliary space: O(n * log(n)).