'''
keys: 
Solutions:
Similar: 
T:
S:
'''
from typing import List
import collections


class Solution:
    # general kSum problem
    # divide k arrays into two groups; O(n^(k/2)) for S and T
    def fourSumCount2(self, A: List[int], B: List[int], C: List[int], D: List[int]) -> int:
        m = {}

        def nSumCount(lists: List[List[int]]) -> int:
            addToHash(lists, 0, 0)
            return countComplements(lists, len(lists) // 2, 0)

        # looping through lists to accumulate
        def addToHash(lists: List[List[int]], i: int, sum: int) -> None:
            if i == len(lists) // 2:
                m[sum] = m.get(sum, 0) + 1
            else: # nested for loop, i.e., choosing one element from lists[0]
                # then choosing another element from lists[1], etc.
                for a in lists[i]:
                    addToHash(lists, i + 1, sum + a)

        # initial i == len(lists)//2, since we are searching the group 2
        def countComplements(lists: List[List[int]], i: int, complement: int) -> int:
            if i == len(lists):
                return m.get(complement, 0)
            cnt = 0
            for a in lists[i]:
                cnt += countComplements(lists, i + 1, complement - a)
            return cnt

        return nSumCount([A, B, C, D])
    
    
    
    # self-made, O(N^2) for S and T
    def fourSumCount(self, A: List[int], B: List[int], C: List[int], D: List[int]) -> int:
        if not any([A, B, C, D]): return 0
        N = len(A)
        memo = collections.defaultdict(int)
        for i in range(N):
            for j in range(N):
                memo[A[i]+B[j]] += 1
        
        res = 0
        for i in range(N):
            for j in range(N):
                if -(C[i] + D[j]) in memo:
                    res += memo[-(C[i] + D[j])]
        return res
    
    # official hashmap
    def fourSumCount1(self, A: List[int], B: List[int], C: List[int], D: List[int]) -> int:
        cnt = 0
        m = {}
        for a in A:
            for b in B:
                m[a + b] = m.get(a + b, 0) + 1
        for c in C:
            for d in D:
                cnt += m.get(-(c + d), 0)
        return cnt
    
sol = Solution()
print (sol.fourSumCount2([1,2], [-2,-1], [-1,2], [0,2]))
        