Problem statement

You are given an array nums of length n. A pivot is an index p with 1 <= p < n that splits the array into two non-empty parts with equal sums: nums[0] + … + nums[p - 1] == nums[p] + … + nums[n - 1].

You may also change at most one element of nums to the value k (or leave the array unchanged).

Return the maximum possible number of pivots after making at most one such change.

Example 1

Input: nums = [2,-1,2], k = 3
Output: 1
Explanation: Change nums[0] to 3: [3,-1,2] has the pivot p = 1 (3 == -1 + 2).

Example 2

Input: nums = [0,0,0], k = 1
Output: 2
Explanation: Leaving the array as is keeps both pivots p = 1 and p = 2.

Example 3

Input: nums = [22,4,-25,-20,-15,15,-16,7,19,-10,0,-13,-14], k = -33
Output: 4
Explanation: Changing nums[2] to -33 gives four pivots.

Constraints

• n == nums.length
• 2 <= n <= 10^5
• -10^5 <= k, nums[i] <= 10^5

Hints for Maximum Number of Ways to Partition an Array

1. Hint 1

   Without any change, pivot p works when 2 * prefix[p] - total == 0. Call d[p] = 2 * prefix[p] - total.

2. Hint 2

   If nums[i] becomes k, let delta = k - nums[i]. Pivots with p <= i keep their left sum and gain delta on the right, so they now need d[p] == delta; pivots with p > i gain delta on the left, so they need d[p] == -delta.

3. Hint 3

   Sweep i from left to right with two hash maps of d values: one for pivots p <= i, one for p > i. Moving i forward moves one pivot from the right map to the left map.

How to solve Maximum Number of Ways to Partition an Array

Measure each pivot by its imbalance d[p] = left - right = 2 * prefix[p] - total. Changing one element shifts the imbalance of every pivot on one side of it by the same amount, so the question becomes "how many pivots on the left have imbalance delta, and how many on the right have imbalance -delta?" — two hash-map lookups per candidate index.

Approach

1. Compute total and d[p] for every pivot p = 1 .. n - 1. Put all of them in a map right (value → count); left starts empty.
2. The answer starts as right[0] — no change at all.
3. For i = 0 .. n - 1: with delta = k - nums[i], the candidate is left[delta] + right[-delta]; update the answer. Then move pivot p = i + 1 (if it exists) from right to left.
4. Return the answer.

Why it works

Changing nums[i] by delta adds delta to the right part of every pivot p <= i (index i lies right of the cut), turning d[p] into d[p] - delta, and adds delta to the left part of every pivot p > i, turning d[p] into d[p] + delta. A pivot is valid iff its new imbalance is zero. At step i, left holds exactly the pivots 1..i and right the pivots i + 1..n - 1, so the two lookups count exactly the pivots valid after the change.

Complexity

• Time — O(n)
• Space — O(n)

Pitfalls

• Sums reach 10^10 in magnitude — use 64-bit integers for prefix sums and imbalances.
• Pivot p has nums[p] in the right part: move it to the left map only after processing i = p - 1.
• Leaving the array unchanged is allowed and must be counted (it can beat every change).

Reference solution

Python

from typing import List
from collections import Counter

def waysToPartition(nums: List[int], k: int) -> int:
    n = len(nums)
    total = sum(nums)
    d = [0] * n
    prefix = 0
    right = Counter()
    for p in range(1, n):
        prefix += nums[p - 1]
        d[p] = 2 * prefix - total
        right[d[p]] += 1
    left = Counter()
    best = right[0]
    for i in range(n):
        delta = k - nums[i]
        best = max(best, left[delta] + right[-delta])
        if i + 1 < n:
            right[d[i + 1]] -= 1
            left[d[i + 1]] += 1
    return best

Java

public static int waysToPartition(int[] nums, int k) {
    int n = nums.length;
    long total = 0;
    for (int x : nums) total += x;
    long[] d = new long[n];
    Map<Long, Integer> right = new HashMap<>(), left = new HashMap<>();
    long prefix = 0;
    for (int p = 1; p < n; p++) {
        prefix += nums[p - 1];
        d[p] = 2 * prefix - total;
        right.merge(d[p], 1, Integer::sum);
    }
    int best = right.getOrDefault(0L, 0);
    for (int i = 0; i < n; i++) {
        long delta = (long) k - nums[i];
        best = Math.max(best, left.getOrDefault(delta, 0) + right.getOrDefault(-delta, 0));
        if (i + 1 < n) {
            right.merge(d[i + 1], -1, Integer::sum);
            left.merge(d[i + 1], 1, Integer::sum);
        }
    }
    return best;
}

C++

int waysToPartition(vector<int>& nums, int k) {
    int n = nums.size();
    long long total = 0;
    for (int x : nums) total += x;
    vector<long long> d(n, 0);
    unordered_map<long long, int> right, left;
    long long prefix = 0;
    for (int p = 1; p < n; p++) {
        prefix += nums[p - 1];
        d[p] = 2 * prefix - total;
        right[d[p]]++;
    }
    int best = right.count(0) ? right[0] : 0;
    for (int i = 0; i < n; i++) {
        long long delta = (long long) k - nums[i];
        int c = 0;
        auto a = left.find(delta);
        if (a != left.end()) c += a->second;
        auto b = right.find(-delta);
        if (b != right.end()) c += b->second;
        best = max(best, c);
        if (i + 1 < n) {
            right[d[i + 1]]--;
            left[d[i + 1]]++;
        }
    }
    return best;
}

JavaScript

var waysToPartition = function(nums, k) {
    var n = nums.length, total = 0;
    for (var t = 0; t < n; t++) total += nums[t];
    var d = new Array(n);
    var right = new Map(), left = new Map();
    var bump = function(m, key, by) { m.set(key, (m.get(key) || 0) + by); };
    var prefix = 0;
    for (var p = 1; p < n; p++) {
        prefix += nums[p - 1];
        d[p] = 2 * prefix - total;
        bump(right, d[p], 1);
    }
    var best = right.get(0) || 0;
    for (var i = 0; i < n; i++) {
        var delta = k - nums[i];
        var c = (left.get(delta) || 0) + (right.get(-delta) || 0);
        if (c > best) best = c;
        if (i + 1 < n) {
            bump(right, d[i + 1], -1);
            bump(left, d[i + 1], 1);
        }
    }
    return best;
};

Also on the editorial tab: C, C#, Go, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

More arrays problems

• Design Movie Rental System Hard
• Prefix and Suffix Search Hard
• Height of Binary Tree After Subtree Removal Queries Hard
• Stamping the Grid Hard
• Smallest Rotation with Highest Score Hard
• Sum of Imbalance Numbers of All Subarrays Hard