CSES Solutions – Array Division

You are given an array arr[] containing N positive integers. Your task is to divide the array into K subarrays so that the maximum sum in a subarray is as small as possible.

Examples:

Input: N = 5, K = 3, arr[] = {2, 4, 7, 3, 5}
Output: 8
Explanation: The 3 subarrays are: {2, 4}, {7} and {3, 5} with sum of 6, 7 and 8 respectively.

Input: N = 4, K = 2, arr = {3, 3, 2, 2}
Output: 6
Explanation: The 2 subarrays are: {3, 3} and {2, 2} with sum of 6 and 4 respectively.

Approach: To solve the problem, follow the below idea:

Let’s assume the minimum of maximum sum of all subarrays to be X. So, for every sum S < X, we know that it is impossible to partition the whole array arr[] into K subarrays such that the sum of each subarray <= S. Also for every sum S >= X, we know that it is always possible to partition the whole array arr[] into K subarrays such that the sum of each subarray <= S. Now, the problem is to find the value of X.

We can find the value of X using Binary Search. If we observe carefully, the minimum possible answer will always be greater than or equal to the maximum element in the array. Also, the maximum possible will always be less than or equal to sum of all elements in the array. Now, we can binary search with low = maximum element in the array and high = total sum of arr[].

We can calculate mid value, and check if we can partition the whole array arr[] into K subarrays such that the sum of each subarray <= mid. If yes, then we can update ans to mid and then reduce our search space by moving high to mid – 1. Else if it is impossible to partition the array arr[] into K subarrays with sum <= mid, then we only reduce our search space by moving low to mid + 1.

We can continue reducing the search space till low <= high, after which we return ans.

Step-by-step algorithm:

• Maintain a function isValid(arr, N, K, mid) to check if we can divide the array arr[] into K subarrays such that each subarray has sum <= mid.
• Maintain a variable ans to store the final answer.
• Initialize low = max element in arr[] and high = sum of all elements in arr[].
• Now, iterate till low <= high
  • Find the mid = (low + high) / 2
  • Check if we can divide the array arr[] into K subarrays such that each subarray has sum <= mid using isValid() function.
  • If it is possible, then update ans = mid and high = mid – 1.
  • Otherwise, update low = mid + 1.
• After all the iterations, return ans as the final answer.

Below is the implementation of the algorithm:

C++

#include <bits/stdc++.h>
#define ll long long int
using namespace std;

// Function to check if we can divide the array arr[] into K
// subarrays such that each subarray has sum <= mid
bool isValid(ll* arr, ll N, ll K, ll mid)
{
    ll partitions = 1, runningSum = 0;
    // find the number of subarrays such that each subarray
    // has sum <= mid
    for (int i = 0; i < N; i++) {
        runningSum += arr[i];
        if (runningSum > mid) {
            runningSum = arr[i];
            partitions += 1;
        }
    }
    // if the number of subarrays is less than or equal to
    // K, then it means that it is possible to divide arr[]
    // into K subarrays with sum of each subarray <= mid
    return partitions <= K;
}

// function to minimize the maximum sum among all subarrays
ll solve(ll* arr, ll N, ll K)
{
    // Define the lower and upper limit of our answer
    ll low = *max_element(arr, arr + N),
    high = accumulate(arr, arr + N, 0LL);
    ll ans = 0;
    // Binary Search to minimize the maximum sum
    while (low <= high) {
        ll mid = (low + high) / 2;
        // If it is possible to divide the array arr[] into
        // K subarrays such that each subarray has sum <=
        // mid, then we update ans and move high to mid-1
        if (isValid(arr, N, K, mid)) {
            ans = mid;
            high = mid - 1;
        }
        // If it is impossible to divide the array arr[]
        // into K subarrays such that each subarray has sum
        // <= mid, move low to mid+1
        else {
            low = mid + 1;
        }
    }
    return ans;
}

int main()
{
    // Sample Input
    ll N = 5, K = 3;
    ll arr[N] = { 2, 4, 7, 3, 5 };

    cout << solve(arr, N, K);
}

Java

import java.util.Arrays;

public class Main {
    // Function to check if we can divide the array arr[] into K
    // subarrays such that each subarray has sum <= mid
    static boolean isValid(long[] arr, int N, int K, long mid) {
        int partitions = 1;
        long runningSum = 0;

        // Find the number of subarrays such that each subarray
        // has sum <= mid
        for (int i = 0; i < N; i++) {
            runningSum += arr[i];
            if (runningSum > mid) {
                runningSum = arr[i];
                partitions += 1;
            }
        }

        // If the number of subarrays is less than or equal to
        // K, then it means that it is possible to divide arr[]
        // into K subarrays with the sum of each subarray <= mid
        return partitions <= K;
    }

    // Function to minimize the maximum sum among all subarrays
    static long solve(long[] arr, int N, int K) {
        // Define the lower and upper limit of our answer
        long low = Arrays.stream(arr).max().getAsLong();
        long high = Arrays.stream(arr).sum();
        long ans = 0;

        // Binary Search to minimize the maximum sum
        while (low <= high) {
            long mid = (low + high) / 2;

            // If it is possible to divide the array arr[] into
            // K subarrays such that each subarray has sum <=
            // mid, then we update ans and move high to mid-1
            if (isValid(arr, N, K, mid)) {
                ans = mid;
                high = mid - 1;
            }
            // If it is impossible to divide the array arr[]
            // into K subarrays such that each subarray has sum
            // <= mid, move low to mid+1
            else {
                low = mid + 1;
            }
        }

        return ans;
    }

    public static void main(String[] args) {
        // Sample Input
        int N = 5, K = 3;
        long[] arr = {2, 4, 7, 3, 5};

        System.out.println(solve(arr, N, K));
    }
}
// This code is contributed by rambabuguphka

JavaScript

// Function to check if we can divide the array arr[] into K
// subarrays such that each subarray has sum <= mid
function isValid(arr, N, K, mid) {
    let partitions = 1;
    let runningSum = 0;
    // find the number of subarrays such that each subarray
    // has sum <= mid
    for (let i = 0; i < N; i++) {
        runningSum += arr[i];
        if (runningSum > mid) {
            runningSum = arr[i];
            partitions++;
        }
    }
    // if the number of subarrays is less than or equal to
    // K, then it means that it is possible to divide arr[]
    // into K subarrays with sum of each subarray <= mid
    return partitions <= K;
}

// function to minimize the maximum sum among all subarrays
function solve(arr, N, K) {
    // Define the lower and upper limit of our answer
    let low = Math.max(...arr);
    let high = arr.reduce((acc, curr) => acc + curr, 0);
    let ans = 0;
    // Binary Search to minimize the maximum sum
    while (low <= high) {
        let mid = Math.floor((low + high) / 2);
        // If it is possible to divide the array arr[] into
        // K subarrays such that each subarray has sum <=
        // mid, then we update ans and move high to mid-1
        if (isValid(arr, N, K, mid)) {
            ans = mid;
            high = mid - 1;
        }
        // If it is impossible to divide the array arr[]
        // into K subarrays such that each subarray has sum
        // <= mid, move low to mid+1
        else {
            low = mid + 1;
        }
    }
    return ans;
}

// Sample Input
let N = 5, K = 3;
let arr = [2, 4, 7, 3, 5];

console.log(solve(arr, N, K));

// This code is contributed by Ayush Mishra

Python3

# Function to check if we can divide the array arr[] into K
# subarrays such that each subarray has sum <= mid
def is_valid(arr, N, K, mid):
    partitions = 1
    running_sum = 0
    # find the number of subarrays such that each subarray
    # has sum <= mid
    for i in range(N):
        running_sum += arr[i]
        if running_sum > mid:
            running_sum = arr[i]
            partitions += 1
    # if the number of subarrays is less than or equal to
    # K, then it means that it is possible to divide arr[]
    # into K subarrays with sum of each subarray <= mid
    return partitions <= K

# function to minimize the maximum sum among all subarrays
def solve(arr, N, K):
    # Define the lower and upper limit of our answer
    low = max(arr)
    high = sum(arr)
    ans = 0
    # Binary Search to minimize the maximum sum
    while low <= high:
        mid = (low + high) // 2
        # If it is possible to divide the array arr[] into
        # K subarrays such that each subarray has sum <=
        # mid, then we update ans and move high to mid-1
        if is_valid(arr, N, K, mid):
            ans = mid
            high = mid - 1
        # If it is impossible to divide the array arr[]
        # into K subarrays such that each subarray has sum
        # <= mid, move low to mid+1
        else:
            low = mid + 1
    return ans

# Sample Input
N = 5
K = 3
arr = [2, 4, 7, 3, 5]

print(solve(arr, N, K))

Output

8

Time Complexity: O(N * logN), where N is the difference between the sum of all elements and the maximum element in the array.
Auxiliary Space: O(1)