In [1]:

import numpy as np
from itertools import combinations, product
from collections import defaultdict

algorithm¶

In [2]:

def sum4(data):
    # store 2-sums
    sum_of_2 = defaultdict(list)
    for i, j in combinations(range(len(data)), 2):
        k = data[i] + data[j]
        sum_of_2[k].append((i, j))

    # match pairs of 2-sums
    sum_of_4 = set()
    for k in sum_of_2:
        if k >= 0 and -k in sum_of_2:
            for i, j in product(sum_of_2[k], sum_of_2[-k]):
                index = tuple(sorted(set(i + j)))
                if len(index) == 4:
                    sum_of_4.add(index)

    return sum_of_4

run¶

In [3]:

n = 10
data = np.random.randint(-n, n, n)
data

Out[3]:

array([  5,   1,  -9,  -6,   4,  -4,   7,   1, -10,  -6])

In [4]:

for index in sum4(data):
    print(index, data[list(index)])

(0, 1, 4, 8) [  5   1   4 -10]
(1, 2, 6, 7) [ 1 -9  7  1]
(1, 3, 4, 7) [ 1 -6  4  1]
(0, 3, 6, 9) [ 5 -6  7 -6]
(0, 4, 7, 8) [  5   4   1 -10]
(1, 4, 7, 9) [ 1  4  1 -6]

In [ ]: