Problem

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 _themaximum possible number of ways to partition _nums to satisfy both conditions after changingat most one element.

Examples

Example 1

1
2
3
4
5

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

1
2
3
4
5
6

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

1
2
3
4

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 <= 10^5
  • -10^5 <= k, nums[i] <= 10^5

Solution

Method 1 – Prefix Sums and Hash Map

Intuition

We want to count the number of ways to split the array into two parts with equal sum, possibly after changing one element to k. We can precompute prefix sums and use hash maps to efficiently count valid partitions before and after a change.

Approach

  1. Compute the total sum of the array and prefix sums for all indices.
  2. For the original array, count the number of valid pivots where the left and right sums are equal.
  3. For each index, consider changing nums[i] to k:
    • The total sum changes by k - nums[i].
    • Use prefix sum counts before and after i to count valid pivots if the change is made at i.
    • Use two hash maps: one for prefix sums to the left, one for the right, and update as you iterate.
  4. Track the maximum number of valid partitions found.
  5. Return the maximum.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

class Solution {
public:
    int waysToPartition(vector<int>& nums, int k) {
        int n = nums.size();
        long long total = 0;
        for (int x : nums) total += x;
        vector<long long> pre(n);
        pre[0] = nums[0];
        for (int i = 1; i < n; ++i) pre[i] = pre[i-1] + nums[i];
        unordered_map<long long, int> left, right;
        for (int i = 0; i < n-1; ++i) right[pre[i]]++;
        int ans = right[total/2];
        for (int i = 0; i < n; ++i) {
            long long diff = k - nums[i];
            long long new_total = total + diff;
            if (new_total % 2 != 0) continue;
            long long half = new_total / 2;
            int cnt = 0;
            if (right.count(half - diff)) cnt += right[half - diff];
            if (left.count(half)) cnt += left[half];
            ans = max(ans, cnt);
            if (i < n-1) {
                right[pre[i]]--;
                left[pre[i]]++;
            }
        }
        return ans;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

func waysToPartition(nums []int, k int) int {
    n := len(nums)
    total := 0
    for _, x := range nums {
        total += x
    }
    pre := make([]int, n)
    pre[0] = nums[0]
    for i := 1; i < n; i++ {
        pre[i] = pre[i-1] + nums[i]
    }
    left := map[int]int{}
    right := map[int]int{}
    for i := 0; i < n-1; i++ {
        right[pre[i]]++
    }
    ans := right[total/2]
    for i := 0; i < n; i++ {
        diff := k - nums[i]
        newTotal := total + diff
        if newTotal%2 != 0 {
            if i < n-1 {
                right[pre[i]]--
                left[pre[i]]++
            }
            continue
        }
        half := newTotal / 2
        cnt := 0
        if v, ok := right[half-diff]; ok {
            cnt += v
        }
        if v, ok := left[half]; ok {
            cnt += v
        }
        if cnt > ans {
            ans = cnt
        }
        if i < n-1 {
            right[pre[i]]--
            left[pre[i]]++
        }
    }
    return ans
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

class Solution {
    public int waysToPartition(int[] nums, int k) {
        int n = nums.length;
        long total = 0;
        for (int x : nums) total += x;
        long[] pre = new long[n];
        pre[0] = nums[0];
        for (int i = 1; i < n; i++) pre[i] = pre[i-1] + nums[i];
        Map<Long, Integer> left = new HashMap<>(), right = new HashMap<>();
        for (int i = 0; i < n-1; i++) right.put(pre[i], right.getOrDefault(pre[i], 0) + 1);
        int ans = right.getOrDefault(total/2, 0);
        for (int i = 0; i < n; i++) {
            long diff = k - nums[i];
            long newTotal = total + diff;
            if (newTotal % 2 != 0) {
                if (i < n-1) {
                    right.put(pre[i], right.get(pre[i]) - 1);
                    left.put(pre[i], left.getOrDefault(pre[i], 0) + 1);
                }
                continue;
            }
            long half = newTotal / 2;
            int cnt = 0;
            cnt += right.getOrDefault(half - diff, 0);
            cnt += left.getOrDefault(half, 0);
            ans = Math.max(ans, cnt);
            if (i < n-1) {
                right.put(pre[i], right.get(pre[i]) - 1);
                left.put(pre[i], left.getOrDefault(pre[i], 0) + 1);
            }
        }
        return ans;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

class Solution {
    fun waysToPartition(nums: IntArray, k: Int): Int {
        val n = nums.size
        val total = nums.sum().toLong()
        val pre = LongArray(n)
        pre[0] = nums[0].toLong()
        for (i in 1 until n) pre[i] = pre[i-1] + nums[i]
        val left = mutableMapOf<Long, Int>()
        val right = mutableMapOf<Long, Int>()
        for (i in 0 until n-1) right[pre[i]] = right.getOrDefault(pre[i], 0) + 1
        var ans = right.getOrDefault(total/2, 0)
        for (i in 0 until n) {
            val diff = k - nums[i]
            val newTotal = total + diff
            if (newTotal % 2L != 0L) {
                if (i < n-1) {
                    right[pre[i]] = right[pre[i]]!! - 1
                    left[pre[i]] = left.getOrDefault(pre[i], 0) + 1
                }
                continue
            }
            val half = newTotal / 2
            var cnt = 0
            cnt += right.getOrDefault(half - diff, 0)
            cnt += left.getOrDefault(half, 0)
            ans = maxOf(ans, cnt)
            if (i < n-1) {
                right[pre[i]] = right[pre[i]]!! - 1
                left[pre[i]] = left.getOrDefault(pre[i], 0) + 1
            }
        }
        return ans
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

class Solution:
    def waysToPartition(self, nums: list[int], k: int) -> int:
        n = len(nums)
        total = sum(nums)
        pre = [0] * n
        pre[0] = nums[0]
        for i in range(1, n):
            pre[i] = pre[i-1] + nums[i]
        from collections import Counter
        left = Counter()
        right = Counter(pre[:n-1])
        ans = right[total // 2] if total % 2 == 0 else 0
        for i in range(n):
            diff = k - nums[i]
            new_total = total + diff
            if new_total % 2 != 0:
                if i < n-1:
                    right[pre[i]] -= 1
                    left[pre[i]] += 1
                continue
            half = new_total // 2
            cnt = right[half - diff] + left[half]
            ans = max(ans, cnt)
            if i < n-1:
                right[pre[i]] -= 1
                left[pre[i]] += 1
        return ans
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

impl Solution {
    pub fn ways_to_partition(nums: Vec<i32>, k: i32) -> i32 {
        use std::collections::HashMap;
        let n = nums.len();
        let total: i64 = nums.iter().map(|&x| x as i64).sum();
        let mut pre = vec![0i64; n];
        pre[0] = nums[0] as i64;
        for i in 1..n {
            pre[i] = pre[i-1] + nums[i] as i64;
        }
        let mut left = HashMap::new();
        let mut right = HashMap::new();
        for i in 0..n-1 {
            *right.entry(pre[i]).or_insert(0) += 1;
        }
        let mut ans = if total % 2 == 0 { *right.get(&(total/2)).unwrap_or(&0) } else { 0 };
        for i in 0..n {
            let diff = k as i64 - nums[i] as i64;
            let new_total = total + diff;
            if new_total % 2 != 0 {
                if i < n-1 {
                    *right.entry(pre[i]).or_insert(0) -= 1;
                    *left.entry(pre[i]).or_insert(0) += 1;
                }
                continue;
            }
            let half = new_total / 2;
            let cnt = right.get(&(half - diff)).unwrap_or(&0) + left.get(&half).unwrap_or(&0);
            ans = ans.max(*cnt);
            if i < n-1 {
                *right.entry(pre[i]).or_insert(0) -= 1;
                *left.entry(pre[i]).or_insert(0) += 1;
            }
        }
        ans as i32
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

class Solution {
    waysToPartition(nums: number[], k: number): number {
        const n = nums.length;
        const total = nums.reduce((a, b) => a + b, 0);
        const pre: number[] = new Array(n);
        pre[0] = nums[0];
        for (let i = 1; i < n; i++) pre[i] = pre[i-1] + nums[i];
        const left = new Map<number, number>();
        const right = new Map<number, number>();
        for (let i = 0; i < n-1; i++) right.set(pre[i], (right.get(pre[i]) ?? 0) + 1);
        let ans = total % 2 === 0 ? (right.get(total/2) ?? 0) : 0;
        for (let i = 0; i < n; i++) {
            const diff = k - nums[i];
            const newTotal = total + diff;
            if (newTotal % 2 !== 0) {
                if (i < n-1) {
                    right.set(pre[i], (right.get(pre[i]) ?? 0) - 1);
                    left.set(pre[i], (left.get(pre[i]) ?? 0) + 1);
                }
                continue;
            }
            const half = newTotal / 2;
            let cnt = (right.get(half - diff) ?? 0) + (left.get(half) ?? 0);
            ans = Math.max(ans, cnt);
            if (i < n-1) {
                right.set(pre[i], (right.get(pre[i]) ?? 0) - 1);
                left.set(pre[i], (left.get(pre[i]) ?? 0) + 1);
            }
        }
        return ans;
    }
}

Complexity

  • ⏰ Time complexity: O(n) — We scan the array and update hash maps in linear time.
  • 🧺 Space complexity: O(n) — For prefix sums and hash maps.