CF 1213D2 - Equalizing by Division (hard version)

Rating: 1600
Tags: brute force, math, sortings
Solve time: 1m 46s
Verified: yes

Solution

Problem Understanding

We are given an array of integers, and we can repeatedly pick any number and divide it by two, rounding down. The goal is to find the minimum number of such operations required to make at least $k$ numbers in the array equal. Each operation decreases a number but never increases it, and the final numbers can include zeros. The input specifies $n$, the number of elements, and $k$, the count of numbers we want equal. The output is a single integer representing the minimal number of division operations.

The constraints are significant: $n$ can reach 200,000, and numbers in the array can also reach 200,000. A naive brute-force approach that tries all sequences of operations for every number would perform on the order of $O(200{,}000^2)$ operations, which is far beyond the time limit. We need an approach that scales linearly or nearly linearly with $n$ and the range of numbers.

Edge cases include situations where many numbers are already equal, where all numbers are distinct powers of two, or where $k$ is equal to $n$. For example, if $n = 3$, $k = 3$, and the array is [8, 4, 2], the minimal operations would be 3: divide 8 to 4 (1), divide 8 to 2 (2), divide 4 to 2 (3) to get three 2s. A careless greedy algorithm that stops at the first equal pair would miss the global minimum.

Approaches

The brute-force solution is conceptually simple. For each number, you could simulate all sequences of divisions, record all the intermediate values, and then try to pick $k$ numbers that can be equal to each intermediate value. You would sum the operations and pick the minimum. This works because every sequence of divisions is predictable, but it is too slow. If every number can take up to 18 steps to reach zero (since 200,000 → 1 is about 18 divisions), then for $n = 200{,}000$, you are already looking at roughly 3.6 million intermediate values per number, leading to over 700 billion steps for a naive enumeration.

The key observation is that each number can be represented as a sequence of values it becomes after each division. For example, 20 → 10 → 5 → 2 → 1 → 0. We do not need to consider which number we divide first; we only need to know, for every possible final value, how many operations each original number needs to reach it. Once we gather these counts, for every candidate final value, we can select the $k$ smallest operation counts to make at least $k$ numbers equal to that value.

This observation allows a much faster approach. We process each number, build a mapping from each reachable value to the number of operations required to reach it, and then, for each candidate value, sort the operations and sum the smallest $k$. This reduces the problem to something like $O(n \log n)$ time overall because each number has at most $\log_2(\text{max a})$ transformations and sorting a small list of size at most $n$ per candidate value is efficient.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(n * max(a) * k)	O(n * max(a))	Too slow
Optimal	O(n log(max(a)) + n log n)	O(n log(max(a)))	Accepted

Algorithm Walkthrough

Initialize a dictionary ops_map where keys are integers representing potential final values, and values are lists of operation counts needed to reach that key.
For each number x in the array, simulate repeated division by two until x reaches zero. Keep a counter op starting at zero for operations. For each value v encountered, append op to ops_map[v] and increment op. This step builds all possible targets and how costly it is for each number to reach them.
After processing all numbers, iterate over every key in ops_map. Each list associated with a key represents the operation counts for numbers that can become this value. Sort the list in ascending order.
For each key, if the list has at least $k$ elements, sum the first $k$ values to compute the total operations needed to make $k$ numbers equal to this value. Track the minimum total operations across all keys.
Output the minimum total operations. This guarantees the minimal number of divisions required to get at least $k$ equal numbers.

Why it works: The algorithm enumerates all values that can be reached from each number and the cost to reach them. By summing the smallest $k$ costs for each candidate value, we guarantee that the chosen $k$ numbers can all become equal to that value with minimal operations. Sorting ensures we always take the least costly numbers, preserving optimality.

Python Solution

import sys
input = sys.stdin.readline
from collections import defaultdict

def main():
    n, k = map(int, input().split())
    a = list(map(int, input().split()))
    
    ops_map = defaultdict(list)
    
    for x in a:
        val = x
        op = 0
        while True:
            ops_map[val].append(op)
            if val == 0:
                break
            val //= 2
            op += 1
    
    ans = float('inf')
    for val, ops_list in ops_map.items():
        if len(ops_list) >= k:
            ops_list.sort()
            total_ops = sum(ops_list[:k])
            if total_ops < ans:
                ans = total_ops
    
    print(ans)

if __name__ == "__main__":
    main()

The code initializes ops_map as a defaultdict of lists to store the cost of reaching each value. The inner loop records all intermediate numbers obtained by dividing by two, along with the number of operations required. The final loop filters candidate values that are achievable by at least $k$ numbers, sorts the operation counts, sums the smallest $k$, and updates the minimal answer. The val == 0 condition ensures we include zero as a reachable value.

Worked Examples

Sample 1

Input:

5 3
1 2 2 4 5

Number	Divisions	Ops per value
1	1 → 0	0 for 1, 1 for 0
2	2 → 1 → 0	0 for 2, 1 for 1, 2 for 0
2	same as above	0 for 2, 1 for 1, 2 for 0
4	4 → 2 → 1 → 0	0 for 4, 1 for 2, 2 for 1, 3 for 0
5	5 → 2 → 1 → 0	0 for 5, 1 for 2, 2 for 1, 3 for 0

Candidate value 2 has operation costs [0, 0, 1, 1]. Sum smallest 3: 0 + 0 + 1 = 1. Answer is 1.

Custom Input

4 4
8 4 2 1

Divisions trace shows all can reach 1:

Number	Ops to reach 1
8	3
4	2
2	1
1	0

Sum for 4 numbers: 0 + 1 + 2 + 3 = 6. Answer is 6.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n log max(a) + n log n)	Each number has at most log(max(a)) divisions; sorting lists of size ≤ n per candidate value dominates
Space	O(n log max(a))	Each intermediate value per number stored; at most n log(max(a)) entries

Given n ≤ 200,000 and max(a) ≤ 200,000, log(max(a)) ≈ 18, total steps are manageable within 2s.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    from __main__ import main
    sys.stdout = io.StringIO()
    main()
    return sys.stdout.getvalue().strip()

# Provided samples
assert run("5 3\n1 2 2 4 5\n") == "1", "sample 1"

# Custom tests
assert run("4 4\n8 4 2 1\n") == "6", "all numbers reduced to 1"
assert run("3 2\n10 20 40\n") == "2", "choose 10->5,20->10 or 20->10,40->20"
assert run("5 1\n1 2 3 4 5\n") == "0", "k=1, no operations needed"
assert run("2 2\n1 1\n") == "0", "all equal already"
assert run("6 3\n16 8 4 2 1 1\n") == "2", "mix of