Scale Engineer

Maximum Number of Ways to Partition an Array

Hard
#1844Time: O(n^2), where `n` is the length of the array. The outer loop runs `n+1` times (for each element to change, plus the no-change case), and inside, counting partitions takes O(n) time.Space: O(n) if a copy of the array is made in each iteration. If the array is modified in-place and then restored, the space complexity is O(1) besides the input storage.
Patterns
CountingEnumerationPrefix Sum
Data structures
ArrayHash Table

Prompt

You are given a 0-indexed integer array nums of length n. The number of ways to partition nums is the number of pivot indices that satisfy both conditions:

  • 1 <= pivot < n
  • nums[0] + nums[1] + ... + nums[pivot - 1] == nums[pivot] + nums[pivot + 1] + ... + nums[n - 1]

You are also given an integer k. You can choose to change the value of one element of nums to k, or to leave the array unchanged.

Return the maximum possible number of ways to partition nums to satisfy both conditions after changing at most one element.

 

Example 1:

Input: nums = [2,-1,2], k = 3
Output: 1
Explanation: One optimal approach is to change nums[0] to k. The array becomes [3,-1,2].
There is one way to partition the array:
- For pivot = 2, we have the partition [3,-1 | 2]: 3 + -1 == 2.

Example 2:

Input: nums = [0,0,0], k = 1
Output: 2
Explanation: The optimal approach is to leave the array unchanged.
There are two ways to partition the array:
- For pivot = 1, we have the partition [0 | 0,0]: 0 == 0 + 0.
- For pivot = 2, we have the partition [0,0 | 0]: 0 + 0 == 0.

Example 3:

Input: nums = [22,4,-25,-20,-15,15,-16,7,19,-10,0,-13,-14], k = -33
Output: 4
Explanation: One optimal approach is to change nums[2] to k. The array becomes [22,4,-33,-20,-15,15,-16,7,19,-10,0,-13,-14].
There are four ways to partition the array.

 

Constraints:

  • n == nums.length
  • 2 <= n <= 105
  • -105 <= k, nums[i] <= 105

Approaches

2 approaches with complexity analysis and trade-offs.

This approach simulates the process directly. It first calculates the number of partitions for the original, unchanged array. Then, it iterates through every possible index i from 0 to n-1, hypothetically changing nums[i] to k. For each change, it recalculates the total sum and all prefix sums, and then counts the number of valid partitions. The maximum count found across all these scenarios (including the no-change case) is the answer.

Algorithm

  • Define a helper function countPartitions(array) that takes an array, calculates its total sum and prefix sums, and returns the number of valid partitions. This takes O(n) time.
  • Calculate the initial number of partitions for the unchanged nums array using the helper function. Store this as max_partitions.
  • Iterate through each index i from 0 to n-1.
  • In each iteration, create a temporary copy of nums, change the element at index i to k.
  • Call countPartitions on the temporary array.
  • Update max_partitions = max(max_partitions, new_count).
  • Return max_partitions.

Walkthrough

Start by calculating the number of ways to partition the original array. This serves as our initial maximum. The condition for a valid partition at pivot p is sum(nums[0...p-1]) == sum(nums[p...n-1]). This can be checked efficiently using prefix sums. Let S be the total sum, and prefix[p] be the sum of the first p elements. The condition is prefix[p] == S - prefix[p], or 2 * prefix[p] == S. We can pre-calculate all prefix sums of the original array in O(n) time. Then, loop through each index i from 0 to n-1. For each i, consider changing nums[i] to k. Inside this loop, create a temporary modified array. Recalculate the prefix sums and the total sum for this modified array. Count the number of valid partitions for the modified array by iterating through all possible pivots p from 1 to n-1 and checking the partition condition. Keep track of the maximum number of partitions found so far. After checking all possible single-element changes, the maximum value recorded is the result.

class Solution {    private int countPartitions(int[] arr) {        int n = arr.length;        long totalSum = 0;        for (int x : arr) {            totalSum += x;        }         if (totalSum % 2 != 0 && n > 0) {            // Optimization: if total sum is odd, no integer prefix sum can be half of it.            // But wait, prefix sums can be non-integers if array has non-integers. Here they are integers.            // The logic `currentPrefixSum * 2 == totalSum` handles this correctly anyway.        }         long currentPrefixSum = 0;        int count = 0;        // Pivot p is between index p-1 and p. 1 <= p < n.        for (int p = 1; p < n; p++) {            currentPrefixSum += arr[p - 1];            if (currentPrefixSum * 2 == totalSum) {                count++;            }        }        return count;    }     public int maxWaysToPartition(int[] nums, int k) {        int n = nums.length;         // Case 1: No change        int maxPartitions = countPartitions(nums);         // Case 2: Change one element        for (int i = 0; i < n; i++) {            int originalValue = nums[i];            nums[i] = k; // Modify the array            maxPartitions = Math.max(maxPartitions, countPartitions(nums));            nums[i] = originalValue; // Backtrack        }         return maxPartitions;    }}

Complexity

Time

O(n^2), where `n` is the length of the array. The outer loop runs `n+1` times (for each element to change, plus the no-change case), and inside, counting partitions takes O(n) time.

Space

O(n) if a copy of the array is made in each iteration. If the array is modified in-place and then restored, the space complexity is O(1) besides the input storage.

Trade-offs

Pros

  • Simple to understand and implement.

  • Directly follows the problem statement.

Cons

  • Inefficient due to repeated calculations. For each of the n possible changes, it performs O(n) work, leading to a quadratic time complexity which is too slow for the given constraints.

Solutions

class Solution {public  int waysToPartition(int[] nums, int k) {    int n = nums.length;    int[] s = new int[n];    s[0] = nums[0];    Map<Integer, Integer> right = new HashMap<>();    for (int i = 0; i < n - 1; ++i) {      right.merge(s[i], 1, Integer : : sum);      s[i + 1] = s[i] + nums[i + 1];    }    int ans = 0;    if (s[n - 1] % 2 == 0) {      ans = right.getOrDefault(s[n - 1] / 2, 0);    }    Map<Integer, Integer> left = new HashMap<>();    for (int i = 0; i < n; ++i) {      int d = k - nums[i];      if ((s[n - 1] + d) % 2 == 0) {        int t = left.getOrDefault((s[n - 1] + d) / 2, 0) +                right.getOrDefault((s[n - 1] - d) / 2, 0);        ans = Math.max(ans, t);      }      left.merge(s[i], 1, Integer : : sum);      right.merge(s[i], -1, Integer : : sum);    }    return ans;  }}

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next →

Next problem