Problem statement

https://leetcode.com/problems/maximum-number-of-ways-to-partition-an-array/

Solution

Actually, not very difficult problem, but I made 6 attempts on contests, because I messed up with indexes. The idea is the following:

If we have numbers [a0, a1, a2, a3, a4, a5, a5, a7, a8, a9] and we want to change say a3 to k so now we have [a0, a1, a2, k, a4, a5, a6, a7, a8, a9].

Then for every partition either left part is from original nums or right part. This is the clue idea which solve this problem.

So, all we need to do is to evaluate cumulative sums both from left and form right and then check how many numbers we can change to k such that sums are equal. We evaluate our cand = (acc1[-1] - nums[i] + k)/2, this is what we want to have in our left part, because if left = right, then left = all/2 and all changed by nums[i] - k. Now we need to understand how many sums are good partitions: for this we keep dictionary d1, which for every cumulatime sum keep list of indexes. We are interested only in indexes which are <= i - 1, because we want to have left part unchanged. Very similar logic is for the cases where right part unchanged, but here we need to use index n - i - 2, because we start from the last element. Finally we update our ans.

Complexity

Time complexity is O(n log n), because for each of n candidates we do 2 binary searches, each of which is O(log n). Space complexity is O(n).

Code

from bisect import bisect

class Solution:
    def waysToPartition(self, nums, k):
        n = len(nums)
        acc1 = list(accumulate(nums))
        d1 = defaultdict(list)
        for idx, num in enumerate(acc1):
            d1[num].append(idx)

        acc2 = list(accumulate(nums[::-1]))
        d2 = defaultdict(list)
        for idx, num in enumerate(acc2):
            d2[num].append(idx)

        ans, ans2 = 0, 0

        for i in range(n - 1):
            if acc1[-1] == 2*acc1[i]: ans2 += 1

        for i in range(0, n):
            cand = (acc1[-1] - nums[i] + k)/2
            lft = bisect(d1[cand], i-1)
            rgh = bisect(d2[cand], n - i - 2)
            ans = max(ans, lft + rgh)

        return max(ans, ans2)