Find 3 elements in sorted array that sum up to 0 (a+b+c = 0)

Question

Is there any better solution than the obvious O(n^3)? I can use one element multiple times so for an array {-1, 0, 2} there is a solution: (-1, -1, 2) or (0, 0, 0).

Answer 1

The classic 3SUM problem can be solved in O(n^2). This can also be solved with O(n^2)

# assume arr is sorted and doesn't contain duplicates

out = []
s = set(arr)
for i in range(len(arr)):
    for j in range(i, len(arr)):
        a = - (arr[i] + arr[j])
        if a >= arr[j] and a in s:
            out.append((arr[i], arr[j], a))