My problem is that we have an array of $N$ integers $(N <=5000)$ on the interval $[-10^6,10^6]$. We also have $Q$ queries $(Q <= 10^5)$ giving us some range in the array.

For each query, we want to find the number of triplets of indices in the range such that the sum of those array values equal to zero. More formally we want to find the number of ways we can choose unordered distinct triples $i, j$ and $k$(within the given query range) such that $a[i]+a[j]+a[k] = 0$.

I'm thinking of doing an $O(N^2)$ pre-computation and $O(\log N)$ time for each query but I am unable to come up with a concrete working idea. Any help would be appreciated.

Edit: queries can be processed offline, as the array does not need to be updated.

  • 1
    $\begingroup$ Can you edit your question to credit the original source where you encountered this task? Thanks! $\endgroup$
    – D.W.
    Commented Jan 19, 2020 at 4:30

1 Answer 1


Let us call the input array $A$. We build the solution in three steps:

First, for each value $v$ between $-10^6$ and $10^6$, build a list with all the indices in $i \in \{1, \dots N\}$, where $A_i = v$ and keep all these lists sorted. One way to achieve this is to iterate over $A$ and append the index of each element to its corresponding list. Let $L_v$ be the list corresponding to a value $v \in \{-10^6, \dots 10^6\}$.

In the second step, we will build a two-dimensional array $C[N][N]$, where \begin{equation} C[i][j] = \begin{cases} |\{k: i+j+k = 0, k \in \{i, \dots j\}\}| &;i\leq j,\\ 0 &;\text{Otherwise}. \end{cases} \end{equation} We can compute $C$ by iterating over all pairs $i, j \in \{1, \dots N\}$, $j\geq i$, and for each such pair, let $u = A_i$ and $v = A_j$. Let $w = -(u + v)$. Note that $u + v + w = 0$, Using binary-search twice, find in $L_w$ the lowest index not less than $i$ and the highest index not greater than $j$. The difference of the indices of these two values plus 1 is the number of elements $w$ between $A_i$ and $A_j$, such that $A_i + A_j + w = 0$, which is exactly $C[i][j]$.

Now as a last pre-processing step, Let us compute the two-dimensional array $S$, the prefix-sums over $C$, that means: $$ S[i][j] = C[i][j] - C[i-1][j] - C[i][j-1] + C[i-1][j-1], $$ where we consider $C[i][j] = 0$ for $i = 0$ or $j = 0$ (since we considered 1-indexed arrays).

Now for each query $i, j$, output $$S[j][j] - S[i-1][j] - S[j][i-1] + S[i-1][i-1],$$ where again, we consider $S[i][j] = 0$ for $i = 0$ or $j=0$.

Based on the previous definitions, try to see why this yields the answer you are looking for.

The running time is $O(N^2\log N)$ pre-processing and $O(1)$ per query. The space complexity is $O(N^2 + M)$, where $M$ is the largest value (equal to $10^6$ in your case), which is needed to keep track of the start of each list $L_i$ for each $i \in [-M, M]$.

  • 1
    $\begingroup$ Shouldn't it be S[j][j]−S[i−1][j]−S[j][i−1]+S[i][i], instead of S[j][j]−S[i−1][j]−S[i][j−1]+S[i][i]. I think this makes sense if you think about it in terms of inclusion-exclusion. $\endgroup$ Commented Jan 19, 2020 at 17:36
  • $\begingroup$ Oh I just corrected two small mistakes. I was thinking about PIE but I wrote it a bit in hurry, thanks for the note $\endgroup$ Commented Jan 19, 2020 at 17:50

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.