CF 1416B - Make Them Equal

Rating: 2000
Tags: constructive algorithms, greedy, math
Solve time: 1m 58s
Verified: no

Solution

Problem Understanding

We are given an array of positive integers and allowed to redistribute values between pairs of elements using a linear operation that depends on the index of the first element. Specifically, we can choose indices $i$ and $j$ and a non-negative integer $x$ and move $x \cdot i$ units from $a_i$ to $a_j$. The goal is to perform at most $3n$ of these operations to make all elements equal. The problem asks either to construct such a sequence or report impossibility.

The key constraint is that every operation must keep all array elements non-negative. Also, the array size $n$ can reach up to $10^4$ and there can be $10^4$ test cases, with the sum of $n$ across all test cases bounded by $10^4$. This ensures that the algorithm must be linear or near-linear in $n$ per test case, ruling out any quadratic redistribution simulation. The individual values $a_i$ can reach $10^5$, so integer overflows are not a concern in Python, but careful arithmetic is necessary to ensure operations respect the index-based multiplication.

Edge cases include arrays where all elements are already equal, where elements cannot be evenly redistributed due to indivisibility constraints, and where the first element is small compared to later elements, making some redistribution impossible if one attempts to naively transfer arbitrary amounts.

For example, an input array [1, 2, 3, 4, 5, 6] cannot be balanced because the sum $21$ is not divisible by $6$, and operations moving multiples of indices cannot fix that. A naive algorithm attempting to just transfer units without checking divisibility could incorrectly produce negative numbers or fail to detect impossibility.

Approaches

A brute-force approach would try to move values in small increments between all pairs until the array becomes equal or we exhaust the $3n$ operations. This is correct in principle because each operation preserves the total sum, but it becomes infeasible for large $n$ since the number of moves could be in the order of the sum of all elements, which can reach $10^9$.

The key insight is that we can structure the operations to move excess values systematically towards the first element, and then redistribute the target value from the first element to every other element. Specifically, we can make each $a_i$ divisible by its index by moving enough from $a_1$, then transfer the remainder of each $a_i$ to $a_1$, effectively consolidating the array's total sum at a single location. Finally, we spread the exact average back to all elements. This deterministic procedure guarantees at most $3n$ operations: one to make divisible, one to move to $a_1$, and one to distribute to all elements.

This structure works because the operation is linear in the index, so we can always solve for the multiplier $x$ to make the remainder zero modulo $i$. By reducing the problem to first making elements divisible and then redistributing, we transform a potentially exponential search into a linear sequence of operations.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(n * max(a_i))	O(n)	Too slow
Greedy Redistribution	O(n)	O(n)	Accepted

Algorithm Walkthrough

Compute the sum $S$ of the array. If $S$ is not divisible by $n$, print $-1$ because equalization is impossible.
Initialize an empty list of operations.
For each index $i$ from $2$ to $n$, check the remainder of $a_i$ modulo $i$. If it is not zero, compute $x = i - (a_i % i)$ and perform the operation $1 \to i$ with this $x$, moving $x$ units from $a_1$ to $a_i$ so that $a_i$ becomes divisible by $i$. Record this operation.
For each $i$ from $2$ to $n$, now that $a_i$ is divisible by $i$, compute $x = a_i // i$ and perform the operation $i \to 1$ with this $x$, transferring the entire divisible portion to $a_1$. Record this operation.
At this point, $a_1$ contains the total sum, and all other $a_i$ are zero. Compute the target value $t = S // n$. For each $i$ from $2$ to $n$, perform the operation $1 \to i$ with $x = t$ to distribute the final equal value to each element. Record each operation.
Print the number of operations and the operations themselves.

This procedure guarantees all elements equal $t$ and uses at most $3n$ operations: one for divisibility adjustment, one for consolidation, and one for distribution per element.

Why it works: the operations preserve the total sum and allow exact control of divisible units due to the index-multiplied transfer. Each element is either increased or decreased in integer multiples of its index, ensuring that the consolidation and redistribution steps do not produce fractional or negative values. The modulo step guarantees that no element becomes negative when making it divisible.

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(-1)
            continue
        ops = []
        target = total // n
        # Step 1: Make all a[i] divisible by i using a[0]
        for i in range(1, n):
            rem = a[i] % (i + 1)
            if rem != 0:
                x = (i + 1) - rem
                a[0] -= x
                a[i] += x
                ops.append((1, i + 1, x))
            # Now transfer a[i] to a[0]
            x = a[i] // (i + 1)
            a[i] -= x * (i + 1)
            a[0] += x * (i + 1)
            if x > 0:
                ops.append((i + 1, 1, x))
        # Step 2: Distribute target from a[0]
        for i in range(1, n):
            ops.append((1, i + 1, target))
        print(len(ops))
        for op in ops:
            print(*op)

if __name__ == "__main__":
    solve()

The solution first checks divisibility to detect impossible cases, then makes each $a_i$ divisible by its index using transfers from $a_1$. After that, it consolidates all divisible amounts at $a_1$, then spreads the target value to all elements. The choice of $x$ in both steps ensures no element becomes negative. The final loop always distributes the average, respecting the operation constraints.

Worked Examples

Example 1: Input [2, 16, 4, 18]. Sum is 40, target is 10.

Step	a	Operation
Initial	[2,16,4,18]	-
Make divisible	[2,16,4,18] → [2,16,4,18] (adjust a_2: 16%2=0, a_3:4%3=1, add 2 from a_1)	1→3:2
Transfer to a_1	[0,16,6,18] → [6,16,0,18]	3→1:2
Transfer to a_1	[6,16,0,18] → [24,0,0,18]	2→1:8
Distribute target	[24,0,0,18] → [10,10,10,10]	1→2:10, 1→3:10, 1→4:10

All elements equal 10.

Example 2: Input [1,2,3,4,5,6]. Sum is 21, not divisible by 6. Output -1. Demonstrates impossibility detection.

These examples confirm that divisibility adjustments and consolidation preserve array invariants and the algorithm correctly detects impossible cases.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n) per test case	Each step iterates over n elements, and each operation is constant time.
Space	O(n)	Array storage plus list of operations, each of size O(n).

Given that the sum of $n$ over all test cases is at most $10^4$, total operations fit comfortably within the 1-second time 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("3\n4\n2 16 4 18\n6\n1 2 3 4