PROBLEM LINK:
Editorialist: Rounaq Jhunjhunu Wala
DIFFICULTY:
EASY-MEDIUM
PREREQUISITES:
Modulo
PROBLEM:
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
EXPLANATION:
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.