4 sum challenge (part 2)

This is a continued discussion from (4 sum challenge) by return count only.


Problem


Given four lists A, B, C, D of integer values, compute how many tuples
(i, j, k, l) there are such that A[i] + B[j] + C[k] + D[l] is zero.


To make problem a bit easier, all A, B, C, D have same length of N
where \$0 \le N \le 500\$. All integers are in the range of \$-2^{28}\$ to \$2^{28} - 1\$
and the result is guaranteed to be at most \$2^{31} - 1\$.


Example:


Input:

A = [ 1, 2]
B = [-2,-1]
C = [-1, 2]
D = [ 0, 2]

Output:

2

Explanation:


The two tuples are:

• (0, 0, 0, 1) -> A[0] + B[0] + C[0] + D[1] = 1 + (-2) + (-1) + 2 = 0

• (1, 1, 0, 0) -> A[1] + B[1] + C[0] + D[0] = 2 + (-1) + (-1) + 0 = 0

I'm wondering if there are any ideas to have a solution less than \$O(n^2)\$ time complexity.

Source code in Python 2.7,

from collections import defaultdict
def four_sum(A, B, C, D):
    sum_map = defaultdict(int)
    result = 0
    for i in A:
        for j in B:
            sum_map[i+j] += 1
    for i in C:
        for j in D:
            if -(i+j) in sum_map:
                result += sum_map[-(i+j)]

    return result

if __name__ == "__main__":
    A = [1, 2]
    B = [-2, -1]
    C = [-1, 2]
    D = [0, 2]
    print four_sum(A,B,C,D)

I’m not sure I can get it any better than \$O(n^2)\$, but I would point you towards the itertools module as a way to tidy up your code by reducing the number of nested for loops. Like so:

def four_sum(A, B, C, D):
    sums = defaultdict(int)
    result = 0
    for (a, b) in itertools.product(A, B):
        sums[a + b] += 1
    for (c, d) in itertools.product(C, D):
        result += sums[-(c + d)]
    return result

I’ve also tweaked the names of some of the variables to make the code easier to follow, and changed it so it only has to compute (c+d) once.

StackExchange Code Review Q#153271, answer score: 7