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

1
# Runtime: 3346 ms (Top 89.52%) | Memory: 46.9 MB (Top 30.64%)
2
class Solution:
3
    def waysToPartition(self, nums: List[int], k: int) -> int:
4
        prefix_sums = list(accumulate(nums))
5
        total_sum = prefix_sums[-1]
6
        best = 0
7
        if total_sum % 2 == 0:
8
            best = prefix_sums[:-1].count(total_sum // 2)  # If no change
9

10
        after_counts = Counter(
11
            total_sum - 2 * prefix_sum for prefix_sum in prefix_sums[:-1]
12
        )
13
        before_counts = Counter()
14

15
        best = max(best, after_counts[k - nums[0]])  # If we change first num
16

17
        for prefix, x in zip(prefix_sums, nums[1:]):
18
            gap = total_sum - 2 * prefix
19
            after_counts[gap] -= 1
20
            before_counts[gap] += 1
21

22
            best = max(best, after_counts[k - x] + before_counts[x - k])
23

24
        return best