CF 1630B - Range and Partition

Rating: 1800
Tags: binary search, constructive algorithms, data structures, greedy, two pointers
Solve time: 3m 26s
Verified: no

Solution

Problem Understanding

We are given an array of integers, and the task is to find a numeric range [x, y] and split the array into exactly k contiguous subarrays so that in each subarray, more than half of the elements fall inside the chosen range. The goal is to minimize the width of the range, y - x. Conceptually, we are trying to "cover" the array with k subarrays where the majority of each subarray lies inside a small numeric window.

The constraints are tight: n can be up to 2·10^5 per test case, and the total sum of n across all test cases is also 2·10^5. This rules out brute-force approaches that examine all possible ranges or all possible subarrays because a naive solution would be O(n^2) or O(n^3), which is infeasible. We need an algorithm roughly O(n log n) or O(n) per test case.

Edge cases arise when the array has many repeated elements, or when k equals 1 or n. For example, if all elements are equal, the range [x, y] is trivial, but a careless approach might attempt to split subarrays incorrectly. If k = n, each subarray contains a single element, so the range must include every element. Another subtle case occurs when the optimal range includes elements that are not contiguous in the array, forcing careful partitioning.

Approaches

The brute-force approach would examine all possible [x, y] ranges, count how many elements in the array fall inside and outside the range for all possible subarrays, and check if a valid split into k subarrays exists. This is correct in principle, but it is O(n^3) because there are O(n^2) ranges and O(n) cost per subarray check. It quickly becomes infeasible for n = 2·10^5.

The key insight is that we do not need to consider every subarray individually to choose the range. We can instead analyze the frequencies of numbers. Let freq[v] be the number of times value v appears in the array. To satisfy the majority condition, the total number of elements inside the range must be at least (n + k) // 2. This comes from the fact that if each subarray has more elements inside than outside, summing across all subarrays gives us a global lower bound on elements in the range.

With this insight, we can efficiently choose the minimal width range [x, y] by using prefix sums on frequency counts of the array values. Once [x, y] is fixed, constructing the subarrays is straightforward with a greedy one-pass: we maintain a balance that increments for elements inside the range and decrements for elements outside. Whenever the balance is positive and we have not yet created k-1 subarrays, we cut the current subarray. The final subarray goes to the end of the array. This guarantees each subarray has more elements inside the range than outside.

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

Algorithm Walkthrough

Count the frequency of each number in the array. We need to know how many times each potential number appears to efficiently evaluate candidate ranges.
Compute prefix sums of these frequencies. Let pref[v] be the number of elements in the array less than or equal to v. Then for any range [x, y], the total number of elements in this range is pref[y] - pref[x-1].
Determine the minimum number of elements needed inside the range to satisfy the majority condition across all k subarrays. This is (n + k) // 2. This formula comes from requiring that each subarray has more inside than outside elements. If we sum the minimal majority from all k subarrays, the total inside elements must be at least (n + k) // 2.
Slide a window [x, y] over the possible numbers using the prefix sums. Start with x = 1 and move y to the smallest value such that the total elements inside [x, y] is at least (n + k) // 2. Record the width y - x. Move x forward and repeat to find the minimal width.
Once the optimal [x, y] is found, construct the subarrays greedily. Iterate over the array and maintain a running balance. For each element, add 1 if it is inside [x, y] and subtract 1 if it is outside. When the balance becomes positive and fewer than k-1 subarrays are created, cut the current subarray and reset the balance. The last subarray ends at the end of the array.

Why it works: By choosing the minimal range with at least (n + k) // 2 elements, we ensure there are enough elements to form k subarrays where the majority condition holds. The greedy balance method guarantees that each subarray individually satisfies the inside > outside requirement.

Python Solution

import sys
input = sys.stdin.readline

def solve():
    t = int(input())
    for _ in range(t):
        n, k = map(int, input().split())
        a = list(map(int, input().split()))
        
        count = [0] * (n + 2)
        for num in a:
            count[num] += 1
        
        pref = [0] * (n + 2)
        for i in range(1, n + 2):
            pref[i] = pref[i-1] + count[i]
        
        needed = (n + k) // 2
        x, y = 1, n
        best_width = n + 1
        left = 1
        for right in range(1, n + 1):
            while pref[right] - pref[left - 1] >= needed:
                if right - left < best_width:
                    best_width = right - left
                    x, y = left, right
                left += 1
        
        print(x, y)
        segments = []
        balance = 0
        last = 0
        cuts = 0
        for i in range(n):
            if x <= a[i] <= y:
                balance += 1
            else:
                balance -= 1
            if balance > 0 and cuts < k - 1:
                segments.append((last + 1, i + 1))
                last = i + 1
                balance = 0
                cuts += 1
        segments.append((last + 1, n))
        for l, r in segments:
            print(l, r)

if __name__ == "__main__":
    solve()

The first part computes frequencies and prefix sums to efficiently find the minimal width range [x, y] that contains at least (n + k) // 2 elements. The second part constructs k subarrays using a running balance to ensure each subarray has a majority of elements inside the range. Boundary handling is crucial: we stop cutting at k-1 subarrays so that the last subarray stretches to the array's end.

Worked Examples

Sample Input 1

4 2
1 2 2 2

i	a[i]	balance	cuts	segments
0	1	-1	0	[]
1	2	0	0	[]
2	2	1	0	[(1,3)]
3	2	1	1	[(1,3)]
Final segment: (4,4)

This confirms the algorithm correctly splits into k=2 subarrays where each has more inside elements than outside.

Sample Input 2

11 3
5 5 5 1 5 5 1 5 5 5 1

i	a[i]	balance	cuts	segments
0-2	5	+3	0	[]
Cut at i=2 -> (1,3)
i=3-6: balance fluctuates, cut at i=6 -> (4,7)
i=7-10: balance positive -> final segment (8,11)

This demonstrates the balance-based greedy cut works with repeated and interleaved elements.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Counting, prefix sum, sliding window, and constructing subarrays each take O(n)
Space	O(n)	Frequency array, prefix sum, and storing segments use O(n)

This fits within the constraints: sum of n ≤ 2·10^5 and 2s time limit.

Test Cases

import sys, io

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

# provided samples
assert run("3\n2 1\n1 2\n4 2\n1 2 2 2\n11 3\n5 5 5 1 5 5