Candy Party (Hard Version) (Codeforces 1869D2)
Tags: Math, bitmasks

This is the hard version of the problem. The only difference is that in this version everyone must give candies to no more than one person and receive candies from no more than one person. Note that a submission cannot pass both versions of the problem at the same time. You can make hacks only if both versions of the problem are solved.After Zhongkao examination, Daniel and his friends are going to have a party. Everyone will come with some candies.There will be $$$n$$$ people at the party. Initially, the $$$i$$$-th person has $$$a_i$$$ candies. During the party, they will swap their candies. To do this, they will line up in an arbitrary order and everyone will do the following no more than once:  Choose an integer $$$p$$$ ($$$1 \le p \le n$$$) and a non-negative integer $$$x$$$, then give his $$$2^{x}$$$ candies to the $$$p$$$-th person. Note that one cannot give more candies than currently he has (he might receive candies from someone else before) and he cannot give candies to himself. Daniel likes fairness, so he will be happy if and only if everyone receives candies from no more than one person. Meanwhile, his friend Tom likes average, so he will be happy if and only if all the people have the same number of candies after all swaps.Determine whether there exists a way to swap candies, so that both Daniel and Tom will be happy after the swaps.

Input: The first line of input contains a single integer $$$t$$$ ($$$1\le t\le 1000$$$) — the number of test cases. The description of test cases follows.The first line of each test case contains a single integer $$$n$$$ ($$$2\le n\le 2\cdot 10^5$$$) — the number of people at the party.The second line of each test case contains $$$n$$$ integers $$$a_1,a_2,\ldots,a_n$$$ ($$$1\le a_i\le 10^9$$$) — the number of candies each person has.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$2\cdot 10^5$$$.

Output: For each test case, print "Yes" (without quotes) if exists a way to swap candies to make both Daniel and Tom happy, and print "No" (without quotes) otherwise.You can output the answer in any case (upper or lower). For example, the strings "yEs", "yes", "Yes", and "YES" will be recognized as positive responses.

Note: In the first test case, the second person gives $$$1$$$ candy to the first person, then all people have $$$3$$$ candies. In the second test case, the fourth person gives $$$1$$$ candy to the second person, the fifth person gives $$$2$$$ candies to the first person, the third person does nothing. And after the swaps everyone has $$$3$$$ candies.In the third test case, it's impossible for all people to have the same number of candies.In the fourth test case, the two people do not need to do anything.

Sample Inputs:
6
3
2 4 3
5
1 2 3 4 5
6
1 4 7 1 5 4
2
20092043 20092043
12
9 9 8 2 4 4 3 5 1 1 1 1
6
2 12 7 16 11 12

Sample Outputs:
Yes
Yes
No
Yes
No
Yes

Solution (Python):
import sys
input = sys.stdin.readline

def readList():
    return list(map(int, input().split()))

def readInt():
    return int(input())

def readInts():
    return map(int, input().split())

def readStr():
    return input().strip()

def solve():
    (n, arr) = (readInt(), readList())
    cnt = [0] * 40
    (pos, neg) = ([0] * 40, [0] * 40)
    t = sum(arr)
    if t % n:
        return 'NO'
    t //= n
    for v in arr:
        (d, p) = (t - v, 0)
        if d == 0:
            continue
        while pow(2, p) <= abs(d):
            p += 1
        if pow(2, p) == 2 * abs(d):
            if d > 0:
                pos[p - 1] += 1
            else:
                neg[p - 1] += 1
        (num, p2) = (pow(2, p) - abs(d), 0)
        while pow(2, p2) < num:
            p2 += 1
        if pow(2, p2) > num:
            return 'NO'
        if d > 0:
            cnt[p] += 1
            cnt[p2] -= 1
        else:
            cnt[p] -= 1
            cnt[p2] += 1
    for i in range(39, 0, -1):
        if cnt[i] > 0:
            v = min(cnt[i], pos[i - 1])
            cnt[i] -= v
            cnt[i - 1] += 2 * v
        elif cnt[i] < 0:
            v = min(-cnt[i], neg[i - 1])
            cnt[i] += v
            cnt[i - 1] -= 2 * v
    return 'YES' if sum([abs(v) for v in cnt]) == 0 else 'NO'
for _ in range(int(input())):
    print(solve())

Solution (Python):
""""""
from collections import Counter
b = [pow(2, i) for i in range(31)]
rev_b = {x: i for (i, x) in enumerate(b)}
diff_i1_i2 = dict()
for i in range(31):
    for j in range(31):
        if i == j:
            continue
        diff_i1_i2[b[j] - b[i]] = (i, j)

def solve(n, a):
    total = sum(a)
    if total % n != 0:
        return 'No'
    mean = total // n
    diffs = [x - mean for x in a]
    diffs = [x for x in diffs if x != 0]
    for diff in diffs:
        if bin(abs(diff)).count('1') == 1:
            continue
        if diff in diff_i1_i2.keys():
            continue
        return 'No'
    two_pow_diff_cnts = Counter()
    degree_cnts = [0] * 31
    for diff in diffs:
        if bin(abs(diff)).count('1') != 1:
            (i, j) = diff_i1_i2[diff]
            degree_cnts[i] -= 1
            degree_cnts[j] += 1
        else:
            two_pow_diff_cnts[diff] += 1
    carry = 0
    for i in range(31):
        degree_cnts[i] += carry
        carry = 0
        (c1, c2) = (two_pow_diff_cnts[b[i]], two_pow_diff_cnts[-b[i]])
        if abs(degree_cnts[i]) > c1 + c2:
            return 'No'
        if (degree_cnts[i] + c1 - c2) % 2 == 1:
            return 'No'
        carry = (degree_cnts[i] + c1 - c2) // 2
    if carry > 0:
        return 'No'
    return 'Yes'

def main():
    t = int(input())
    allans = []
    for _ in range(t):
        n = int(input())
        a = readIntArr()
        ans = solve(n, a)
        allans.append(ans)
    multiLineArrayPrint(allans)
    return
import sys
input = sys.stdin.buffer.readline

def oneLineArrayPrint(arr):
    print(' '.join([str(x) for x in arr]))

def multiLineArrayPrint(arr):
    print('\n'.join([str(x) for x in arr]))

def multiLineArrayOfArraysPrint(arr):
    print('\n'.join([' '.join([str(x) for x in y]) for y in arr]))

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

def makeArr(defaultValFactory, dimensionArr):
    dv = defaultValFactory
    da = dimensionArr
    if len(da) == 1:
        return [dv() for _ in range(da[0])]
    else:
        return [makeArr(dv, da[1:]) for _ in range(da[0])]

def queryInteractive(a, b):
    print('? {} {}'.format(a, b))
    sys.stdout.flush()
    return int(input())

def answerInteractive(ans):
    print('! {}'.format(ans))
    sys.stdout.flush()
inf = float('inf')
from math import gcd, floor, ceil
import math
for _abc in range(1):
    main()

Solution (Python):
import sys
ipt = sys.stdin.readline

def iin():
    return int(ipt())

def lmin():
    return list(map(int, ipt().split()))
T = 30
d = {2 ** i: i for i in range(T)}

def solve():
    N = iin()
    A = lmin()
    S = sum(A)
    if S % N:
        print('NO')
        return
    M = S // N
    if all((a == M for a in A)):
        print('YES')
        return
    cnt_in1 = [0] * T
    cnt_out1 = [0] * T
    cnt_in2 = [0] * T
    cnt_out2 = [0] * T
    for i in range(N):
        if M - A[i] == 0:
            continue
        if M - A[i] in d:
            cnt_out2[d[M - A[i]]] += 1
        if -M + A[i] in d:
            cnt_in2[d[-M + A[i]]] += 1
        t = 1
        for j in range(T):
            if M - (A[i] - t) in d:
                cnt_in1[j] += 1
                cnt_out1[d[M - (A[i] - t)]] += 1
                break
            t *= 2
        else:
            print('NO')
            return
    for i in range(T - 1, 0, -1):
        if cnt_in1[i] > cnt_out1[i]:
            if cnt_in1[i] - cnt_out1[i] <= cnt_in2[i - 1]:
                cnt_in1[i - 1] += cnt_in1[i] - cnt_out1[i]
                cnt_out1[i - 1] -= cnt_in1[i] - cnt_out1[i]
            else:
                print('NO')
                return
        elif cnt_in1[i] < cnt_out1[i]:
            if cnt_out1[i] - cnt_in1[i] <= cnt_out2[i - 1]:
                cnt_out1[i - 1] += cnt_out1[i] - cnt_in1[i]
                cnt_in1[i - 1] -= cnt_out1[i] - cnt_in1[i]
            else:
                print('NO')
                return
    if cnt_in1[0] == cnt_out1[0]:
        print('YES')
    else:
        print('NO')
for _ in range(iin()):
    solve()