Computer >> Computer tutorials >  >> Programming >> Python

4Sum II in Python

Suppose we have four lists A, B, C, D of integer values, we have to compute how many tuples (i, j, k, l) there are such that A[i] + B[j] + C[k] + D[l] is zero. Consider all A, B, C, D have same length of N where 0 ≤ N ≤ 500. Remember all integers are in the range of -228 to 228 - 1 and the result is guaranteed to be at most 231 - 1. So if the inputs are A = [1, 2], B = [-2, -1], C = [-1, 2], D = [0, 2], then the output will be 2. This is because there are two tuples, they are (0, 0, 0, 1) so A[0] + B[0] + C[0] + D[1] = 1 + (-2) + (-1) + 2 = 0, and another tuple is (1, 1, 0, 0), A[1] + B[1] + C[0] + D[0] = 2 + (-1) + (-1) + 0 = 0

To solve this, we will follow these steps −

  • make one map called sums
  • for each element i in A
    • for each element j in B
      • if i + j is not in the sums map, then set sums[i + j] := 1
      • otherwise increase the value of sums[i + j]
  • counter := 0
  • for each element i in C
    • for each element j in D
      • if (-1)*(i + j) in sums, then counter := counter + sums[-1 * (i + j)]
  • return counter

Example

Let us see the following implementation to get better understanding −

class Solution(object):
   def fourSumCount(self, A, B, C, D):
      sums ={}
      for i in A:
         for j in B:
            if i+j not in sums:
               sums[i+j] = 1
            else:
               sums[i+j] +=1
      counter = 0
      for i in C:
         for j in D:
            if -1 * (i+j) in sums:
               #print(-1 * (i+j))
               counter+=sums[-1*(i+j)]
      return counter
ob1 = Solution()
print(ob1.fourSumCount([1,2], [-2,-1], [-1,2], [0,2]))

Input

[1,2]
[-2,-1]
[-1,2]
[0,2]

Output

2