Editorialist: Rounaq Jhunjhunu Wala
Given a list of numbers A, we need to find the number of triplets (x,y,z) such that x+y+z \equiv 0 \mod 3
The easiest approach is to run three nested loops, which picks one element each and check whether the sum is a multiple of 3 or not.
ans = 0 for i in range 1 to N-2 for j in range i+1 to N-1 for k in range j+1 to N if (A[i]+A[j]+A[k])%3 == 0: ans = ans + 1 return ans
The obvious disadvantage for the above is the time complexity [O(N^3)]. We can do better by using properties of modulus (
(a+b)%m = ((a%m)+(b%m))%m). We first create a new array B such that
B[i] = A[i]%3. Now, we count the frequency of elements in the above list. Let the number of zeros, ones and twos be a,b and c respectively, then we can directly find the answer from a, b and c as follows. The table below show the different combinations of elements that form a valid triplet. Values are \mod 3 and C(a,b) = \binom ab.
a b c total combinations 0 0 0 0 C(a,3) 1 1 1 0 C(b,3) 2 2 2 0 C(c,3) 0 1 2 0 a*b*c
So, we can infer from above that the answer is \binom a3+\binom b3+\binom c3+abc.