Count pairs

Practice

2.9 (17 votes)

Sorting

Hashmap

Implementation

Hash tables

Data structures

Basics of hash tables

Hashing

Math

Observation

Problem

88% Success 5391 Attempts 30 Points 1.5s Time Limit 256MB Memory 1024 KB Max Code

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 $$

Sample Input

2
4 1 2
3 2 1 0
8 0 3
9 2 1 8 4 5 9 10

Sample Output

5
7

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)$$.

Code Editor

You need to be logged in to access the code editor

Please wait while we load the editor

Submissions

Please login to view your submissions

You need to be logged in to access the code editor

Please wait while we load the editor