3SUM problem (finding triplets) in better than O(n^2)

Question

Consider the following problem:

Given an unsorted array of integers, find all triplets that satisfy x^2 + y^2 = z^2.

For example if given array is 1, 3, 7, 5, 4, 12, 13, the answer should be 5, 12, 13 and 3, 4, 5

I suggest the below algorithm with complexity O(n^2):

• Sort the array in descending order, O(nlogn)
• square each element, O(n)

Now it reduces to the problem of finding all triplets(a,b,c) in a sorted array such that a = b+c.

The interviewer was insisting on a solution better than O(n^2).

I have read 3SUM problem on Wikipedia, which emphasizes problem can be solved in O(n+ulogu) if numbers are in range [-u,u] assuming the array can be represented as a bit vector. But I am not able to get a clear picture of further explanations.

Can someone please help me in understanding what is going on with a nice example?

Answer 1

First of all. Finding all triplets in worst case is O(n^3). Suppose you have n=3k numbers. K of them are 3, k are 4 and k are 5.

3,....,3,4,....,4,5,.....5

There are k^3 = n^3/27 = O(n^3) such triplets. Just printing them takes O(n^3) time.

Next will be explaining of 3SUM problem in such form:

There are numbers s_1, ..., s_n each of them in range [-u;u]. How many triplets a,b,c there are that a+b=c?

1. transforming. Get 2*u numbers a_-u, ..., a_0, a_1, ... a_u. a_i is amount of numbers s_j, that s_j = i. This is done in O(n+u)

2. res = a_0 * sum(i=-u..u, i=/=0, C(a_i, 2)) + C(a_0,3) a_0 = 0

3. Build a polynom P(x) = sum(i = -u...u, a_i*x^(i+u).

4. Find Q(x) = P(x)*P(x) using FFT.

5. Note that Q(x) = sum(i=-2u..2u, b_i*x^(i+2*u)), where b_i is number of pairs s_u,s_k that s_u+s_k = i (This includes using same number twice).

6. For all even i do b_i = b_i - a_(i/2). This will remove using same number twice.

7. Sum all b_i*a_i/2 - add this to res.

Example: to be more simple, I will assume that range for numbers is [0..u] and won't use any +u in powers of x. Suppose that we have numbers 1,2,3 - a_0 = 0, a_1 = 1, a_2 = 1, a_3 = 1

• res = 0

• P(x) = x + x^2 + x^3

• Q(x) = x^2 +2x^3+3x^4+2x^5+x^6

• After subtracting b_2 = 0, b_3 = 2, b_4 = 2, b_5 = 2, b_6 = 0

• res += 0*1/2 + 2*1/2 + 2*0/2 + 2*0/2 + 6*0/2 = 1