PRCN16C - Editorial

PROBLEM LINK:

Practice
Contest

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.