Question

Consider there are three array list each of equal length and having positive, negative and zero in them. I had to write a program to find combinations of which sum comes to zero. So basically, if the arrays were:-

A = {0, -1, 2}
B = {0, -1, -2}
C = {0, -2, 0}

O/P: A[0] + B[0] + C[0], A[2] + B[2] + C[2] etc.

I could think of two ways,
1. Have 3 for loops and calculate the sum using a[i] + b[j] + c[k], if zero print the index.
Big O will be O(N^3)
2. Have two for loop but use Binary Search to find the third element which would give the sum as zero.
Big O will be O(N^2LogN)

Any other ways?

Thanks.

EDIT:
Based on the answers given below, my first soln is the fastest possible. But if the question is about “finding” the number of combinations and NOT printing them, then please see Grigor Gevorgyan answer below.

Answer 1

Can be done in O(n^2) with 2 pointers method.

Sort the arrays. Now do following:
Set ans = 0.
Have an external loop running through array a with index i. Now set j = 0, k = n – 1.
Look at sum = a[ i ] + b[ j ] + c[ k ].
If sum < 0, increase j.
If sum > 0 decrease k.
If sum == 0, find the range of elements equal to b[ j ] and c[ k ] and add ranges lengths product to the answer. Then set j and k to first elements out of that ranges.
This works because arrays are sorted and addition is a linear function.
Internal part runs in O(n), with overall O(n^2) complexity.

Example code in C++:

sort( a, a + n );
sort( b, b + n );
sort( c, c + n );
ans = 0;
for( i = 0; i < n; ++i )
{
   j = 0, k = n - 1;
   while( j < n && k > 0 )
   {
      sum = a[ i ] + b[ j ] + c[ k ];
      if( sum < 0 ) ++j;
          else if( sum > 0 ) --k;
          else 
          { 
             // find the equal range
             for( jj = j; jj < n && b[ jj ] == b[ j ]; ++jj );
             for( kk = k; kk >= 0 && c[ kk ] == c[ k ]; --kk );
             // add pairs quantity from these ranges
             ans += ( jj - j ) * ( k - kk );
             j = jj, k = kk;
          }
}

Note: Sorting of array a is not necessary, just did it to look good 🙂