PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Cakewalk
PREREQUISITES:
None
PROBLEM:
Given an array A with elements in [1, 3], count the number of pairs of elements whose sum is a prime.
EXPLANATION:
The array elements are in [1, 3], so there aren’t really that many options for what the sum of two of them can actually be - in fact, the sum must be an integer in the range [2, 6].
The only prime numbers in the range [2, 6] are \{2, 3, 5\}.
So, our aim is to count the number of pairs of elements whose sum equals one of these three values.
There are now two ways to solve the problem:
- Use the fact that N \le 100, so that \mathcal{O}(N^2) is fast enough.
This allows us to simply iterate over each pair of elements, and check if their sum if one of \{2, 3, 5\} or not, solving the problem. - A faster solution is to observe that:
The only way to obtain a sum of 2 is via 1+1.
The only way to obtain a sum of 3 is via 1+2.
The only way to obtain a sum of 3 is via 2+3.
So, if c_1, c_2, c_3 are the counts of 1, 2, 3 in the array respectively,- There are c_1\cdot c_2 ways to choose a pair of (1, 2), whose sum is 3 (c_1 ways to choose a 1, and c_2 ways to choose a 2).
- There are c_2\cdot c_3 ways to choose a pair of (2, 3), whose sum is 5.
- There are \frac{c_1 \cdot (c_1 - 1)}{2} ways to choose a pair of (1, 1) from distinct indices whose sum is 2.
- The answer is hence the sum of these three quantities.
The second solution runs in \mathcal{O}(N) time since it only requires the computation of the values c_1, c_2, c_3.
TIME COMPLEXITY:
\mathcal{O}(N) or \mathcal{O}(N^2) per testcase.
CODE:
Editorialist's code (PyPy3)
for _ in range(int(input())):
n = int(input())
a = list(map(int, input().split()))
freq = [0]*4
for x in a: freq[x] += 1
print(freq[1]*freq[2] + freq[2]*freq[3] + freq[1]*(freq[1]-1)//2)