Question

2025. Maximum Number of Ways to Partition an Array

中文文档

Description

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

Solutions

Solution 1: Prefix Sum + Hash Table

We can preprocess to get the prefix sum array $s$ corresponding to the array $nums$, where $s[i]$ represents the sum of the array $nums[0,…i-1]$. Therefore, the sum of all elements in the array is $s[n - 1]$.

If we do not modify the array $nums$, the condition for the sums of the two subarrays to be equal is that $s[n - 1]$ must be even. If $s[n - 1]$ is even, then we calculate $ans = \frac{right[s[n - 1] / 2]}{2}$.

If we modify the array $nums$, we can enumerate each modification position $i$, change $nums[i]$ to $k$, then the change in the total sum of the array is $d = k - nums[i]$. At this time, the sum of the left part of $i$ remains unchanged, so the legal split must satisfy $s[i] = s[n - 1] + d - s[i]$, that is, $s[i] = \frac{s[n - 1] + d}{2}$. Each prefix sum of the right part has increased by $d$, so the legal split must satisfy $s[i] + d = s[n - 1] + d - (s[i] + d)$, that is, $s[i] = \frac{s[n - 1] - d}{2}$. We use hash tables $left$ and $right$ to record the number of times each prefix sum appears in the left and right parts, respectively. Then we can calculate $ans = max(ans, left[\frac{s[n - 1] + d}{2}]) + right[\frac{s[n - 1] - d}{2}]$.

Finally, we return $ans$.

The time complexity is $O(n)$, and the space complexity is $O(n)$. Here, $n$ is the length of the array $nums$.

class Solution:
  def waysToPartition(self, nums: List[int], k: int) -> int:
    n = len(nums)
    s = [nums[0]] * n
    right = defaultdict(int)
    for i in range(1, n):
      s[i] = s[i - 1] + nums[i]
      right[s[i - 1]] += 1

    ans = 0
    if s[-1] % 2 == 0:
      ans = right[s[-1] // 2]

    left = defaultdict(int)
    for v, x in zip(s, nums):
      d = k - x
      if (s[-1] + d) % 2 == 0:
        t = left[(s[-1] + d) // 2] + right[(s[-1] - d) // 2]
        if ans < t:
          ans = t
      left[v] += 1
      right[v] -= 1
    return ans

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;
  }
}

class Solution {
public:
  int waysToPartition(vector<int>& nums, int k) {
    int n = nums.size();
    long long s[n];
    s[0] = nums[0];
    unordered_map<long long, int> right;
    for (int i = 0; i < n - 1; ++i) {
      right[s[i]]++;
      s[i + 1] = s[i] + nums[i + 1];
    }
    int ans = 0;
    if (s[n - 1] % 2 == 0) {
      ans = right[s[n - 1] / 2];
    }
    unordered_map<long long, int> left;
    for (int i = 0; i < n; ++i) {
      int d = k - nums[i];
      if ((s[n - 1] + d) % 2 == 0) {
        int t = left[(s[n - 1] + d) / 2] + right[(s[n - 1] - d) / 2];
        ans = max(ans, t);
      }
      left[s[i]]++;
      right[s[i]]--;
    }
    return ans;
  }
};

func waysToPartition(nums []int, k int) (ans int) {
  n := len(nums)
  s := make([]int, n)
  s[0] = nums[0]
  right := map[int]int{}
  for i := range nums[:n-1] {
    right[s[i]]++
    s[i+1] = s[i] + nums[i+1]
  }
  if s[n-1]%2 == 0 {
    ans = right[s[n-1]/2]
  }
  left := map[int]int{}
  for i, x := range nums {
    d := k - x
    if (s[n-1]+d)%2 == 0 {
      t := left[(s[n-1]+d)/2] + right[(s[n-1]-d)/2]
      if ans < t {
        ans = t
      }
    }
    left[s[i]]++
    right[s[i]]--
  }
  return
}