PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
Greedy
PROBLEM:
An array of length 3N is called balanced if it contains N occurrences each of 1, 2, 3.
Given K \le 10^7, find any array of length \le 10^4 that contains exactly K balanced subarrays.
EXPLANATION:
Balanced subarrays must have a length that’s a multiple of 3, so an array of length N can have at best (N-2) + (N-5) + (N-8) + \ldots balanced subarrays (since there are N-K+1 subarrays of length K).
This summation is approximately \frac{N^2}{6} (in fact, in general slightly lower than this).
Note that for N = 10^4, \frac{N^2}{6} \approx 1.6 \cdot 10^7, which is only slightly more than the maximum allowed K of 10^7.
This means in general, we definitely need to be able to create arrays where (nearly) every subarray that can be balanced, is balanced.
This is not too hard in and of itself - a simple construction is to just use
[1, 2, 3, 1, 2, 3, 1, 2, 3, \ldots]
It’s easy to see that the periodic nature of this array ensures that any subarray whose length is a multiple of 3 is balanced.
Now, we need exactly K balanced subarrays, so we can’t simply use periodic arrays.
Instead, what we’ll do is find a way to join some of these periodic arrays together while appropriately separating them so that different parts don’t interfere with each others.
Let B_i denote the array formed by i copies of [1, 2, 3].
So for example, B_3 = [1, 2, 3, 1, 2, 3, 1, 2, 3].
It can be verified that B_i will have exactly \frac{i}{2} \cdot (3i - 1) balanced subarrays.
Let S_i = \frac{i}{2} \cdot (3i - 1).
There is now a somewhat natural greedy algorithm that should come to mind:
- Start with an empty array A.
We have K balanced subarrays required. - Repeat the following while K \gt 0:
- Choose the largest integer i such that S_i \le K.
- Append the array B_i to A.
- Reduce K by S_i.
This will surely give us K balanced subarrays - the only issue is, it will likely give us more than K, because we didn’t account for subarrays that cross different “blocks”.
For example, if K = 2, our algorithm will first append [1, 2, 3] to A and decrement K by 1, and then again append [1, 2, 3] to A and decrement K by 1.
However the final array is now [1, 2, 3, 1, 2, 3], which has more than 2 balanced subarrays.
To account for this, we need a way of separating blocks so that no new balanced subarrays will be created.
Let’s analyze what this might look like with the above example of K = 2: we have the two blocks [1, 2, 3] and [1, 2, 3], how can they be separated?
If a 1 is appended to the first block, it immediately creates another balanced subarray of the form [2, 3, 1], so that’s not something we can do.
If a 3 is inserted before the second block, it again creates a balanced subarray of the form [3, 1, 2], so that isn’t allowed either.
That leaves the element 2. However, separating the blocks with a single 2 won’t work either: we end up with [1, 2, \underline{3, 2, 1}, 2, 3], where the underlined segment is a new balanced subarray created.
However, nothing forces us to use a single copy of 2 - we can instead use two copies!
That is, create the array [1, 2, 3, 2, 2, 1, 2, 3].
It can be verified that there are now exactly two balanced subarrays, as we wanted.
This idea works in general too - simply separate any two blocks by [2, 2], and then no subarray that crosses blocks can ever be balanced (because such a subarray will always end up containing more copies of 2 then either 1 or 3).
So for example if our building blocks are [1, 2, 3, 1, 2, 3], [1, 2, 3], [1, 2, 3] then the final array will be
This makes our final algorithm the following:
- Repeat the following while K \gt 0:
- Choose the largest integer i such that S_i \le K.
- Append the array B_i to A.
- Append [2, 2] to A.
- Reduce K by S_i.
We now have a fairly simple construction that achieves K balanced subarrays - the only question is whether it uses \le 10^4 elements always.
Luckily for us, it does!
Proof
Let g(K) denote the length of the array created by our algorithm for when K balanced subarrays are required.
Recall that we defined S_i = \frac{i}{2} \cdot (3i-1).There’s a simple recurrence for this function: g(K) = 3x + 2 + g\left(K - S_x\right), where x is the largest integer such that S_x \le K.
The base case is, of course, g(0) = 0.Let’s compute an upper bound on g(K).
In particular, we can bound the value of the chosen x, by \sqrt K.
To see why: substitute x = \sqrt{K} into the expression for S_x and it results in \frac{3K - \sqrt{K}}{2}, which is always \ge K, because K \ge \sqrt{K} for all K \ge 1.Next, we look at the remaining part.
Note that K - S_x \lt S_{x+1} - S_x, because S_{x+1} \gt K.
Now,S_{x+1} - S_x = \frac{(x+1)\cdot (3(x+1)-1)}{2} - \frac{x\cdot (3x-1)}{2} = \frac{1}{2} \cdot (6x + 2) = 3x + 1This tells us that K - S_x \lt 3\sqrt K + 1.
So, \displaystyle g(K) \le \sqrt K + 2 + \max_{x \le 3\sqrt K + 1} g(x).
The second term is now again bounded by something of the form \displaystyle\sqrt{3\sqrt{K}+1} + 2 + \max_{x \le c} g(x) where c = 3\sqrt{3\sqrt{K}+1} + 1.
In general, we’re just repeatedly taking the square root (with a small multiplier of 3), which will shrink in value very quickly.
Our sum thus looks like\sqrt{K} + \sqrt{3\sqrt{K} + 1} + \sqrt{3\sqrt{3\sqrt{K} + 1} + 1} + \ldotsAs you can see, the terms decay rapidly, so the value is dominated by the first term.
Indeed, by induction it can be shown that this expression is \le 3\sqrt{K} for large enough K (where “large enough” here is like 100 or so).With K = 10^7, 3\sqrt{K} is approximately 9500 so it fits into the limit of 10^4 nicely.
In fact, the factor of 3 is a bit generous - the actual factor is even lower so in the actual worst case our algorithm only uses a length of 8000 or so.Note that we didn’t account for the effect of the +2's, but it’s easy to see that they don’t really matter much: a +2 will be added for each successive function call, but because the length drops by a square-root factor each time, the number of function calls is bounded by \mathcal{O}(\log \log K) so the overall effect is tiny and we’re still quite safe.
Note that because K \le 10^7, it’s also possible to just explicitly verify yourself that the length is small enough for all K, which can be done in linear time using dynamic programming.
In particular, the maximum length we’ll ever create is a bit over 8000.
TIME COMPLEXITY:
\mathcal{O}(\sqrt K) per testcase.
CODE:
Editorialist's code (PyPy3)
for _ in range(int(input())):
k = int(input())
ans = []
for i in reversed(range(1, 10**4)):
while i*(3*i-1)//2 <= k:
for j in range(3*i):
ans.append(j%3 + 1)
ans.append(2)
ans.append(2)
k -= i*(3*i-1)//2
print(len(ans))
print(*ans)