Count pairs | Practice Problems

Count pairs

★★★★★

2.9

17 votes

Sorting, Hashmap, Implementation, Hash Tables, Data Structures, Basics of Hash Tables, Hashing, Math, Observation

Details

Submissions

Discussion

Editorial

Problem

You are given an array $A$ consisting of $N$ non-negative integers. You are also given 2 integers $p$ (a prime number) and $k$ . You are required to count number of pairs $(i, j)$ where, $1 \leq i < j \leq N$ and satisying:

$(A_i^2 + A_j^2 + A_i*A_j) \text{ mod }p = k$

where $a \text{ mod } p = b$ means that $b$ is the remainder when $a$ is divided by $p$ . In particular, $0 \leq b < p$ .

You are given $T$ test cases.

Input format

The first line contains a single integer $T$ denoting the number of test cases.
For each test case:
- The first line contains three space-separated integers $N$ , $k$ , and $p$ denoting the length of the array, the required remainder/modulo, and the prime number respectively.
- The second line contains $N$ space-separated integers denoting the integer array $A$ .

Output format

For each test case, print the number of pairs satisfying the equation in a separate line.

Constraints

$1 \le T \le 1000 \\ 1 \le N \le 5 \times 10^5 \\ 1 \le p \le 10^9 \\ \text{p is prime} \\ 0 \le k < p \\ 0 \le A_i \le 10^9 \\ \text{Sum of N over all test cases does not exceed } 10^6$

Time Limit: 1.5

Memory Limit: 256

Source Limit:

Explanation

In the first sample, $N = 4, k = 1, p = 2, A = [3, 2, 1, 0]$ . Caculations for each pair $(i, j)$ are shown below:

$(1, 2)$ : $A_1^2 + A_2^2 + A_1*A_2 \text{ mod } p = (9 + 4 + 6) \text{ mod } 2 = 1$ .
$(1, 3)$ : $A_1^2 + A_3^2 + A_1*A_3 \text{ mod } p = (9 + 1 + 3) \text{ mod } 2 = 1$ .
$(1, 4)$ : $A_1^2 + A_4^2 + A_1*A_4 \text{ mod } p = (9 + 0 + 0) \text{ mod } 2 = 1$ .
$(2, 3)$ : $A_2^2 + A_3^2 + A_2*A_3 \text{ mod } p = (4 + 1 + 4) \text{ mod } 2 = 1$ .
$(2, 4)$ : $A_2^2 + A_4^2 + A_2*A_4 \text{ mod } p = (4 + 0 + 0) \text{ mod } 2 = 0$ .
$(3, 4)$ : $A_3^2 + A_4^2 + A_3*A_4 \text{ mod } p = (1 + 0 + 0) \text{ mod } 2 = 1$ .
Therefore, the number of pairs satisfying the equation are 5.

In the second sample, $N = 8, k = 0, p = 3, A = [9, 2, 1, 8, 4, 5, 9, 10]$ . Valid pairs $(i, j)$ are: $(1, 7), (2, 4), (2, 6), (3, 5), (3, 8), (4, 6), (5, 8)$ .