Count Pairs (Codeforces 1189E)
Tags: Others, number theory

You are given a prime number $$$p$$$, $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$, and an integer $$$k$$$. Find the number of pairs of indexes $$$(i, j)$$$ ($$$1 \le i < j \le n$$$) for which $$$(a_i + a_j)(a_i^2 + a_j^2) \equiv k \bmod p$$$.

Input: The first line contains integers $$$n, p, k$$$ ($$$2 \le n \le 3 \cdot 10^5$$$, $$$2 \le p \le 10^9$$$, $$$0 \le k \le p-1$$$). $$$p$$$ is guaranteed to be prime.The second line contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$0 \le a_i \le p-1$$$). It is guaranteed that all elements are different.

Output: Output a single integer — answer to the problem.

Note: In the first example:$$$(0+1)(0^2 + 1^2) = 1 \equiv 1 \bmod 3$$$.$$$(0+2)(0^2 + 2^2) = 8 \equiv 2 \bmod 3$$$.$$$(1+2)(1^2 + 2^2) = 15 \equiv 0 \bmod 3$$$.So only $$$1$$$ pair satisfies the condition.In the second example, there are $$$3$$$ such pairs: $$$(1, 5)$$$, $$$(2, 3)$$$, $$$(4, 6)$$$.

Sample Inputs:
3 3 0
0 1 2

6 7 2
1 2 3 4 5 6

Sample Outputs:
1

3

Solution (Python):
def count(n, p, k, L):
    s = {}
    for x in L:
        m = (x ** 4 - x * k) % p
        if m in s:
            s[m] += 1
        else:
            s[m] = 1
    res = 0
    for x in s:
        y = s[x]
        res += y * (y - 1) // 2
    return res
(n, p, k) = [int(x) for x in input().split()]
L = [int(x) for x in input().split()]
print(count(n, p, k, L))

Solution (Python):
(n, p, k) = [int(i) for i in input().split()]
l = [int(i) for i in input().split()]
di = {}
ans = 0
for i in l:
    x = (i * i * i * i - k * i) % p
    if x in di:
        ans += di[x]
        di[x] += 1
    else:
        di[x] = 1
print(ans)

Solution (Python):
def main():
    (n, p, k) = map(int, input().split())
    ans = list(map(int, input().split()))
    for i in range(n):
        temp = ans[i]
        ans[i] = (temp ** 4 - temp * k) % p
    mp = {}
    for i in ans:
        mp.setdefault(i, 0)
        mp[i] = mp[i] + 1
    temp = 0
    for i in mp.values():
        temp = i * (i - 1) // 2 + temp
    print(temp)
if __name__ == '__main__':
    main()