CF 1869D1 - Candy Party (Easy Version)

CF 1869D1 - Candy Party (Easy Version)

Rating: 1700
Tags: bitmasks, constructive algorithms, graphs, implementation, math
Solve time: 1m 28s
Verified: yes

Solution

Problem Understanding

We have a party with n people, each holding a certain number of candies a_i. Each person will give a positive number of candies to exactly one other person, and each person will receive candies from exactly one other person. The candies given must be powers of two, and no one can give more than they currently have.

The goal is twofold: first, each person gives to exactly one and receives from exactly one (a perfect one-to-one mapping of candy transfers), and second, after all transfers, everyone ends up with the same number of candies. The input consists of multiple test cases, each specifying n and the list a. The output is "Yes" if a valid transfer configuration exists and "No" otherwise.

The constraints are tight: n can reach 2×10^5 per test case, with the sum across all test cases also capped at 2×10^5. This rules out any solution that tries all permutations of transfers, which would be factorial in n.

Edge cases include situations where all candies are already equal, which is trivially "Yes", and cases where the sum of candies is not divisible by n, which immediately makes a solution impossible. Small n = 2 is also an edge case because it is the minimum size for a one-to-one swap.

Approaches

The brute-force method is to try every permutation of giving candies from one person to another in powers of two, then simulate the resulting state and check if everyone ends up equal. This would involve iterating over all permutations of size n, each with up to log(max(a_i)) power-of-two splits, leading to complexity O(n! × log(max(a_i))). This is unfeasible for n up to 2×10^5.

The key observation is that every person's net gain or loss must be exactly target - a_i, where target is the final number of candies each person should have. For the average to be equal for all, target must be the total number of candies divided by n. If this is not an integer, a solution is impossible.

Given the power-of-two transfer restriction, we need each person's net gain/loss to be representable as a sum of powers of two. Specifically, the differences between a_i and target can be decomposed into sums of distinct powers of two. Because everyone gives to exactly one person and receives from exactly one person, the set of outgoing and incoming powers-of-two transfers must balance perfectly.

The approach that emerges is to check whether each difference target - a_i can be expressed as a sum of powers of two. By counting how many of each power of two are needed across all positive and negative differences, we can verify whether the total sums of each power of two match exactly. If they do, a valid assignment exists. This avoids simulating every permutation.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(n! × log(max(a_i)))	O(n)	Too slow
Optimal	O(n × log(max(a_i)))	O(log(max(a_i)))	Accepted

Algorithm Walkthrough

1. For each test case, read n and the list a. Compute the total number of candies S and determine target = S // n. If S % n != 0, print "No" and continue, because it is impossible for everyone to end up equal.
2. Initialize two counters: give and receive, each as a map from power-of-two values to their counts.
3. For each person i, calculate diff = a[i] - target. If diff > 0, this person must give away candies. If diff < 0, they must receive -diff candies.
4. For each non-zero diff, decompose it into powers of two. This is done by iterating over bits from 0 to 30 (since a_i ≤ 10^9), incrementing the corresponding power in give if diff > 0 or in receive if diff < 0.
5. After processing all people, for a solution to exist, the multiset of powers-of-two needed to give must match exactly the multiset needed to receive. If they match, print "Yes". Otherwise, print "No".

Why it works: the decomposition into powers of two ensures that each person's net difference can be satisfied by giving or receiving powers-of-two candies. Because everyone must give to exactly one and receive from exactly one person, the total counts for each power must balance. This invariant guarantees correctness without explicitly simulating all transfer sequences.

Python Solution

import sys
input = sys.stdin.readline

def solve():
    t = int(input())
    for _ in range(t):
        n = int(input())
        a = list(map(int, input().split()))
        total = sum(a)
        if total % n != 0:
            print("No")
            continue
        target = total // n
        
        give = {}
        receive = {}
        
        for x in a:
            diff = x - target
            if diff == 0:
                continue
            power = 0
            if diff > 0:
                while diff:
                    if diff & 1:
                        give[power] = give.get(power, 0) + 1
                    diff >>= 1
                    power += 1
            else:
                diff = -diff
                while diff:
                    if diff & 1:
                        receive[power] = receive.get(power, 0) + 1
                    diff >>= 1
                    power += 1
        
        if give == receive:
            print("Yes")
        else:
            print("No")

solve()

The code first calculates the target candies per person. Then it builds the counts of powers of two each person must give or receive. Comparing these counts ensures that every required transfer can be performed without violating the one-to-one mapping constraint. Careful attention is paid to correctly handle negative differences and bitwise decomposition. Using a dictionary avoids creating large arrays unnecessarily.

Worked Examples

Example 1

Input: 3 2 4 3

Target: (2+4+3)//3 = 3

Person	Candies	Diff	Powers-of-two needed
1	2	-1	receive 2^0
2	4	1	give 2^0
3	3	0	none

give = {0:1}, receive = {0:1} → match → "Yes"

Example 2

Input: 6 1 4 7 1 5 4

Target: (1+4+7+1+5+4)//6 = 3

Person	Candies	Diff	Powers-of-two needed
1	1	-2	receive 2^1
2	4	1	give 2^0
3	7	4	give 2^2
4	1	-2	receive 2^1
5	5	2	give 2^1
6	4	1	give 2^0

give = {0:2,1:1,2:1}, receive = {1:2} → no match → "No"

This demonstrates handling multiple powers and mismatched sums.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n × log(max(a_i)))	Each person's diff is decomposed into ≤30 powers of two.
Space	O(log(max(a_i)))	Only need to store counts of each power of two up to 2^30.

With n ≤ 2×10^5 and a_i ≤ 10^9, this fits comfortably in the 2-second limit.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    sys.stdout = io.StringIO()
    solve()
    return sys.stdout.getvalue().strip()

# 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\nNo"

# custom cases
assert run("1\n2\n1 3\n") == "Yes", "n=2 simple swap"
assert run("1\n3\n1 1 1\n") == "Yes", "all equal"
assert run("1\n4\n1 2 3 5\n") == "No", "sum not divisible"
assert run("1\n5\n2 2 2 2