Maximum Number of Ways to Partition an Array LeetCode solution in Java

1
class Solution {
2
  // time - O(n), space - O(n)
3
  public int waysToPartition(int[] nums, int k) {
4
    int n = nums.length;
5

6
    int[] pref = new int[n];
7
    pref[0] = nums[0];
8
    Map<Integer, Integer> count = new HashMap<>(); // contribution of prefixes without changing
9
    count.put(pref[0], 1);
10

11
    for (int i = 1; i < n - 1; i++) {
12
      pref[i] += pref[i - 1] + nums[i];
13
      count.put(pref[i], count.getOrDefault(pref[i], 0) + 1);
14
    }
15
    pref[n - 1] +=
16
        pref[n - 2]
17
            + nums[
18
                n - 1]; // last step doesn't add into 'count' map, because we're trying to find at
19
    // least two parts.
20

21
    int sum = pref[n - 1];
22
    int max = 0;
23
    if (sum % 2 == 0) max = count.getOrDefault(sum / 2, 0); // max without changing
24

25
    Map<Integer, Integer> countPrev =
26
        new HashMap<>(); // contribution of prefixes before current step
27
    for (int i = 0; i < n; i++) {
28
      int diff = k - nums[i];
29
      int changedSum = sum + diff;
30
      if (changedSum % 2 == 0)
31
        max =
32
            Math.max(
33
                max,
34
                count.getOrDefault(changedSum / 2 - diff, 0)
35
                    + countPrev.getOrDefault(changedSum / 2, 0));
36
      count.put(pref[i], count.getOrDefault(pref[i], 0) - 1);
37
      countPrev.put(pref[i], countPrev.getOrDefault(pref[i], 0) + 1);
38
    }
39
    return max;
40
  }
41
}