CF 1869D2 - Candy Party (Hard Version)

CF 1869D2 - Candy Party (Hard Version)

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

Solution

Problem Understanding

We have a party with n people, each bringing a certain number of candies. Each person can give exactly a power-of-two number of candies to another person, and every person can give to at most one person and receive from at most one person. The goal is to determine if it is possible to rearrange the candies so that every person ends up with the same number of candies.

Input provides multiple test cases. For each test case, we get the number of people n and an array a of initial candy counts. The output is "Yes" if it is possible to redistribute candies to make all counts equal under the given rules, otherwise "No".

Constraints require careful consideration. n can be up to 2*10^5, and total n across all test cases does not exceed 2*10^5. Each candy count can be up to 10^9. This rules out algorithms that attempt to simulate all possible transfers directly or try all subsets of giving and receiving - the brute-force would require exponentially many operations. Instead, we need an approach that leverages the properties of powers of two and differences between candy counts.

A non-obvious edge case arises when everyone already has the same number of candies. The answer should be "Yes" even if no moves are necessary. Another subtle case is when differences are not powers of two or cannot be matched one-to-one between donors and receivers. For example, if one person has 1 candy and another has 3, it is impossible to transfer in a single power-of-two step to equalize them because 2 is not a power-of-two difference here.

Approaches

The brute-force approach would attempt to enumerate all possible transfers for each person and check if a valid redistribution exists. For n = 10^5, this approach fails immediately because the number of candidate moves grows exponentially, and checking all combinations is infeasible.

The key insight is that each transfer is a power-of-two and that each person can be part of at most one giving and one receiving action. This restriction allows us to reason about differences rather than simulating transfers. Let target be the average number of candies. Any person with more than target must give away exactly the difference, and anyone with fewer must receive exactly the difference. If all differences can be expressed as a sum of distinct powers of two and can be paired uniquely (each donor to a receiver), then the redistribution is possible.

The greedy approach is as follows: calculate the target candies per person, compute the differences, and check if every surplus and deficit can be resolved by unique power-of-two transfers. This reduces the problem to counting bits - the surplus and deficit in each bit position must balance exactly.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(2^n)	O(n)	Too slow
Bitmask/Greedy	O(n * log(max(a_i)))	O(1)	Accepted

Algorithm Walkthrough

1. For each test case, read n and the array a of candy counts. Compute the sum of candies, total.
2. Check if total is divisible by n. If not, equal distribution is impossible, output "No".
3. Calculate target = total // n.
4. Initialize two lists: give for the amount each person must give (positive differences) and receive for the amount each person must receive (negative differences in absolute value).
5. For each person, calculate diff = a[i] - target. If diff > 0, append to give; if diff < 0, append absolute value to receive.
6. For each value in give and receive, break it down into sums of distinct powers of two using bit operations. Track how many transfers of each power-of-two are needed.
7. Compare the count of required transfers for each power-of-two in give and receive. If they match exactly, output "Yes". Otherwise, output "No".

Why it works: every person gives to at most one person and receives from at most one person, so the only allowed moves are exactly one donor to one receiver for each power-of-two amount. Balancing the counts of powers-of-two ensures that a one-to-one mapping exists for all required transfers.

Python Solution

import sys
input = sys.stdin.readline

def can_equalize(a):
    n = len(a)
    total = sum(a)
    if total % n != 0:
        return False
    target = total // n
    give_bits = [0] * 31
    receive_bits = [0] * 31
    for x in a:
        diff = x - target
        if diff > 0:
            for i in range(31):
                if diff & (1 << i):
                    give_bits[i] += 1
        elif diff < 0:
            diff = -diff
            for i in range(31):
                if diff & (1 << i):
                    receive_bits[i] += 1
    return give_bits == receive_bits

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

The solution breaks each difference into powers of two and counts how many of each bit are required. Using a list of size 31 is sufficient because a_i <= 10^9 (fits in 30 bits). Comparing these counts guarantees that every surplus can match a deficit with a unique one-to-one power-of-two transfer.

Worked Examples

Example 1: a = [2, 4, 3]

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

Differences: [-1, +1, 0]

• Give: 1 (from 4 to 2) → bit 0 = 1
• Receive: 1 (from 2 to 4) → bit 0 = 1

Counts match → Yes

Example 2: a = [1, 4, 7, 1, 5, 4]

Target = (1+4+7+1+5+4)/6 = 22/6 ≈ 3.666 → not integer → No

This demonstrates that the algorithm correctly detects impossible cases when the total is not divisible by n.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n * log(max(a_i)))	For each test case, we process each difference into at most 30 bits
Space	O(1)	Only need two arrays of size 31 to store bit counts

Given n ≤ 2*10^5 total across all test cases and log(max(a_i)) ≤ 30, this solution runs comfortably within 2 seconds.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    output = io.StringIO()
    sys.stdout = output
    exec(open("candy_party.py").read())  # assuming solution saved as candy_party.py
    return output.getvalue().strip()

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 edge cases
assert run("1\n2\n1 1\n") == "Yes", "all equal values"
assert run("1\n3\n1 1 2\n") == "No", "cannot redistribute to equal"
assert run("1\n4\n1 3 3 1\n") == "Yes", "powers of two redistribution possible"
assert run("1\n3\n5 5 5\n") == "Yes", "no move needed"
assert run("1\n2\n1 3\n") == "Yes", "single power-of-two move possible"

Test input	Expected output	What it validates
2\n1 1	Yes	Minimum-size input, all equal
3\n1 1 2	No	Impossible redistribution
4\n1 3 3 1	Yes	Powers-of-two decomposition works
3\n5 5 5	Yes	Already equal, no moves needed
2\n1 3	Yes	Single power-of-two adjustment

Edge Cases

For a = [1, 3]

Target = 2

Differences: [-1, +1]

• Give: 1 → bit 0 = 1
• Receive: 1 → bit 0 = 1

Counts match, so the algorithm correctly returns "Yes".

For a = [1, 1, 2]

Target = 4/3 → not integer, algorithm returns "No" immediately, avoiding unnecessary calculations.

All edge