Leetcode 1013. Partition Array Into Three Parts With Equal Sum

Problem statement

https://leetcode.com/problems/partition-array-into-three-parts-with-equal-sum/

Solution

First we need to check that sum is divisible by 3. Then we need to look at prefixes and suffixes and check if we have values sm//3 among them. Finally, if we have, check that the first place in prefix is smaller than the last place of suffix minus one.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def canThreePartsEqualSum(self, arr):
        sm = sum(arr)
        if sm % 3 != 0: return False
        pref = list(accumulate(arr))
        if sm//3 not in pref: return False
        suff = list(accumulate(arr[::-1]))
        if sm//3 not in suff: return False
        idx1, idx2 = pref.index(sm//3), len(arr) - 1 - suff.index(sm//3)
        return idx1 + 1 < idx2