Delete Two Elements (Codeforces 1598C)
Tags: Math, two pointers

Monocarp has got an array $$$a$$$ consisting of $$$n$$$ integers. Let's denote $$$k$$$ as the mathematic mean of these elements (note that it's possible that $$$k$$$ is not an integer). The mathematic mean of an array of $$$n$$$ elements is the sum of elements divided by the number of these elements (i. e. sum divided by $$$n$$$).Monocarp wants to delete exactly two elements from $$$a$$$ so that the mathematic mean of the remaining $$$(n - 2)$$$ elements is still equal to $$$k$$$.Your task is to calculate the number of pairs of positions $$$[i, j]$$$ ($$$i < j$$$) such that if the elements on these positions are deleted, the mathematic mean of $$$(n - 2)$$$ remaining elements is equal to $$$k$$$ (that is, it is equal to the mathematic mean of $$$n$$$ elements of the original array $$$a$$$).

Input: The first line contains a single integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of testcases.The first line of each testcase contains one integer $$$n$$$ ($$$3 \le n \le 2 \cdot 10^5$$$) — the number of elements in the array.The second line contains a sequence of integers $$$a_1, a_2, \dots, a_n$$$ ($$$0 \le a_i \le 10^{9}$$$), where $$$a_i$$$ is the $$$i$$$-th element of the array.The sum of $$$n$$$ over all testcases doesn't exceed $$$2 \cdot 10^5$$$.

Output: Print one integer — the number of pairs of positions $$$[i, j]$$$ ($$$i < j$$$) such that if the elements on these positions are deleted, the mathematic mean of $$$(n - 2)$$$ remaining elements is equal to $$$k$$$ (that is, it is equal to the mathematic mean of $$$n$$$ elements of the original array $$$a$$$).

Note: In the first example, any pair of elements can be removed since all of them are equal.In the second example, there is no way to delete two elements so the mathematic mean doesn't change.In the third example, it is possible to delete the elements on positions $$$1$$$ and $$$3$$$, or the elements on positions $$$4$$$ and $$$5$$$.

Sample Inputs:
4
4
8 8 8 8
3
50 20 10
5
1 4 7 3 5
7
1 2 3 4 5 6 7

Sample Outputs:
6
0
2
3

Solution (Python):
def solve(n, a):
    s2 = sum(a) * 2
    if s2 % n:
        return 0
    target = s2 // n
    d = {}
    res = 0
    for v in a:
        diff = target - v
        count = d.get(diff, 0)
        res += count
        count = d.get(v, 0)
        d[v] = count + 1
    return res
t = int(input())
for i_t in range(t):
    n = int(input())
    a = list(map(int, input().split(' ')))
    print(solve(n, a))

Solution (Python):
from collections import deque
from collections import OrderedDict
import queue
import heapq
import math
from re import I
import sys
sys.setrecursionlimit(10 ** 6)

def sf():
    return [int(x) for x in input().split(' ')]

def sfi():
    return int(input())

def sfs():
    return input()

def printf(x):
    print(x)
    sys.stdout.flush()

def main():
    t = sfi()
    for tc in range(t):
        fr = {}
        n = sfi()
        ls = sf()
        sm = sum(ls) << 1
        if sm % len(ls):
            print('0')
            continue
        tot = int(sm / len(ls))
        res = 0
        for i in ls:
            x = tot - i
            if x in fr:
                res += fr[x]
            if i in fr:
                fr[i] += 1
            else:
                fr[i] = 1
        print(res)
if __name__ == '__main__':
    main()

Solution (Python):
import sys

def solve(n, a):
    mean = sum(a) / n
    k = sum(a) - mean * (n - 2)
    m = {}
    res = 0
    for i in range(n):
        if k - a[i] in m:
            res += m[k - a[i]]
        m[a[i]] = 1 + m.get(a[i], 0)
    return res
if __name__ == '__main__':
    try:
        t = int(input().strip())
        for _ in range(t):
            n = int(input().strip())
            a = list(map(int, input().strip().split()))
            result = solve(n, a)
            print(result)
    except Exception as e:
        print(f'Exception: {e}')
        sys.exit(1)