LeetCode: 454. 4Sum II

First thing to try is iterating nums1, nums2, nums3 and looking for the corresponding element in nums4. This will result in O(N**3) time complexity, which raises TLE.

class Solution:
    def fourSumCount(self, nums1: List[int], nums2: List[int], nums3: List[int], nums4: List[int]) -> int:
        from collections import Counter
        c = Counter(nums4)
        ret = 0
        for n1 in nums1:
            for n2 in nums2:
                for n3 in nums3:
                    ret += c[-(n1 + n2 + n3)]
        return ret

 

How can we reduce the time complexity?

Making use of the fact that all arrays have same length of N would be great.

The idea is to pair (nums1, nums2) and (nums3, nums4) each, and to store all sum results generated by both pairs in the seperate counters. --> O(N**2)

Then, iterate through all sum results in one counter. --> O(N**2)

class Solution:
    def fourSumCount(self, nums1: List[int], nums2: List[int], nums3: List[int], nums4: List[int]) -> int:
        from collections import Counter
        c1 = Counter()
        c2 = Counter()
        for n1 in nums1:
            for n2 in nums2:
                c1[n1+n2] += 1
        for n3 in nums3:
            for n4 in nums4:
                c2[n3+n4] += 1
        ret = 0
        for k, v in c1.items():
            ret += c2[-k] * v
        return ret