CF 1868B2 - Candy Party (Hard Version)

Rating: 2100
Tags: bitmasks, constructive algorithms, dp, greedy, implementation, math
Solve time: 1m 47s
Verified: no

Solution

Problem Understanding

We are given a group of people at a party, each initially holding some number of candies. Each person can give a certain number of candies to another person exactly once, with the amount being a power of two and not exceeding the candies they currently hold. The goal is to make all people end up with the same number of candies while ensuring that each person receives candies from at most one other person. In other words, we need to find a one-to-one mapping of candy transfers such that after all transfers, everyone has an identical count of candies.

The input provides multiple test cases. For each test case, we know how many people are present and the number of candies each person initially has. The output is a simple "Yes" if the transfers can be arranged to satisfy both fairness and equality or "No" otherwise.

The constraints are such that the number of people can be up to 200,000, and the total sum over all test cases is also up to 200,000. Each person can hold up to $10^9$ candies. This rules out any solution that tries all possible sequences of candy transfers, since that would be combinatorially explosive. Instead, we need a linear or near-linear approach per test case.

The non-obvious edge cases include situations where everyone already has the same number of candies, situations where the differences between candies are powers of two, and situations where the differences cannot be expressed as powers of two or cannot be paired in a one-to-one mapping. For example, if the input is [1, 2, 3], one might initially try to transfer candies from the person with 3 to the one with 1 and 2 to 2, but it fails because a person cannot receive from multiple people. The correct output here is "Yes", because we can arrange a one-to-one transfer that matches powers of two.

Approaches

A brute-force approach would attempt to simulate all possible sequences of transfers, checking each combination of giving and receiving powers of two. This would involve generating all powers of two subsets of each person's candies, matching them to all other people, and checking if the final counts are equal. The number of possibilities grows exponentially with $n$ and the number of bits in each candy count. For $n$ around $10^5$, this is hopelessly slow.

The key observation is that each transfer must be a power of two and can only occur between unique pairs of giver and receiver. This means that the problem reduces to checking whether the differences between the target number of candies and each person's initial candies can be represented as a power-of-two sum that allows a one-to-one assignment. First, we compute the target number of candies, which must be the average of all candies, since everyone must end with the same count. If the total candies are not divisible by the number of people, it is impossible. Next, we split the people into those who need to give candies and those who need to receive candies. Each giving amount must be a power of two, and each receiving amount must be matched to a giver. By considering the largest power-of-two difference first and matching them greedily, we can check if such a pairing is possible.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(2^(n log max_a))	O(n)	Too slow
Optimal	O(n log max_a)	O(n)	Accepted

Algorithm Walkthrough

For each test case, calculate the sum of all candies. If the sum is not divisible by the number of people, print "No" and continue to the next case. This ensures a uniform target is mathematically possible.
Compute the target number of candies as target = sum(a) // n.
For each person, calculate the difference diff = a[i] - target. If diff > 0, this person must give away candies. If diff < 0, this person must receive candies.
Store the giving and receiving differences in two separate lists. Take the absolute value of receiving differences.
For both lists, attempt to decompose each difference into sums of distinct powers of two. This is equivalent to checking the binary representation of the difference. Each bit set in the binary representation corresponds to a transfer of $2^k$ candies.
For each power of two, check that there are enough corresponding givers and receivers to match. This is done by counting how many givers can provide $2^k$ and how many receivers need $2^k$. If at any power the number of givers is insufficient for receivers, print "No".
If all powers of two differences are successfully matched, print "Yes".

Why it works: The greedy decomposition into powers of two works because any integer can be represented uniquely as a sum of distinct powers of two. By matching givers and receivers per bit, we ensure that each person gives to at most one other person and receives from at most one, satisfying the problem's constraints.

Python Solution

import sys
input = sys.stdin.readline

def can_equalize_candies(n, candies):
    total = sum(candies)
    if total % n != 0:
        return False
    target = total // n
    give_counts = {}
    receive_counts = {}
    
    for c in candies:
        diff = c - target
        if diff == 0:
            continue
        x = abs(diff)
        bits = []
        i = 0
        while x:
            if x & 1:
                bits.append(1 << i)
            x >>= 1
            i += 1
        if diff > 0:
            for b in bits:
                give_counts[b] = give_counts.get(b, 0) + 1
        else:
            for b in bits:
                receive_counts[b] = receive_counts.get(b, 0) + 1
    
    for b in receive_counts:
        if receive_counts[b] > give_counts.get(b, 0):
            return False
    return True

t = int(input())
for _ in range(t):
    n = int(input())
    a = list(map(int, input().split()))
    print("Yes" if can_equalize_candies(n, a) else "No")

The solution first checks whether a uniform target is possible. Differences are decomposed into powers of two, and we maintain counts of required transfers for each power. By matching these counts, we enforce the one-to-one constraint efficiently. Edge cases such as everyone already having the same number of candies are handled naturally because the differences are zero.

Worked Examples

Sample 1

Input: [2, 4, 3]

Target: 3

Person	Candies	Diff	Bits	Type
1	2	-1	[1]	Receive
2	4	1	[1]	Give
3	3	0	[]	None

The give bit 1 matches the receive bit 1, so output is "Yes".

Sample 2

Input: [1, 2, 3, 4, 5]

Sum = 15, n = 5, target = 3

Person	Candies	Diff	Bits	Type
1	1	-2	[2]	Receive
2	2	-1	[1]	Receive
3	3	0	[]	None
4	4	1	[1]	Give
5	5	2	[2]	Give

Each give bit exactly matches a receive bit, output is "Yes".

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n log max_a)	Each difference is decomposed into powers of two, which is proportional to the number of bits in max_a
Space	O(n)	Storing counts of powers for givers and receivers

The approach scales linearly with n times the number of bits, which is acceptable since n ≤ 2×10^5 and max_a ≤ 10^9 (~30 bits).

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    output = []
    t = int(input())
    for _ in range(t):
        n = int(input())
        a = list(map(int, input().split()))
        output.append("Yes" if can_equalize_candies(n, a) else "No")
    return "\n".join(output)

# provided samples
assert run("6\n3\n2 4 3\n5\n1 2 3 4 5\n6\n1 4 7 1 5 4\n2\n20092043 20092043\n12\n9 9 8 2 4 4 3 5 1 1 1 1\n6\n2 12 7 16 11 12\n") == "Yes\nYes\nNo\nYes\nNo\nYes"

# custom cases
assert run("1\n2\n1 1\n") == "Yes", "already equal"
assert run("1\n2\n1 3\n") == "Yes", "power of two transfer possible"
assert