PROBLEM LINK:
Author: Богдан Пастущак
Tester: Suchan Park
DIFFICULTY:
Medium
PREREQUISITES:
Dynamic programming and bitsets.
PROBLEM:
You are given two positive integers: N <= 10^6 and K <= 4. Consider numbers 1^K, 2^K, 3^K, …, N^K. You need to divide them in two groups (such that every number is present in exactly one group) with as close sum of elements as possible. Also you need to print one such partition.
QUICK EXPLANATION:
- Let’s call this two groups A and B. Also let’s define by Sum(G) as sum of elements of group G. We are interested in such partition, that minimizes answer = |Sum(A) - Sum(B)|.
- It is easy to notice that answer must be odd if N is equal 1 or 2 modulo 4 and even otherwise.
- What if N is small? We can do standard knapsack-like boolean dynamic programming dp[i][j] - is it possible to distribute numbers from 1^K, 2^K, 3^K, … , i^K between two groups A[i] and B[i], such that Sum(A[i]) - Sum(B[i]) = j? We can optimize this with bitsets easily. With such dp we can also restore group division.
- Now, what if N is too big? One may notice that if N is big enough, it is always possible to divide number in two groups in such way that answer is either 0 or 1. But how to restore group division based on this? It turns out that the following is holding: for any positive integer x it is always possible to divide numbers x^K, (x + 1)^K, (x + 2)^K, … , (x + 2^{(K + 1)} - 1)^K in two groups with equal sum of elements. Based on this fact we can step by step distribute 2^{(K + 1)} biggest numbers, making N quite small, and then apply dynamic programming from the previous step.
EXPLANATION:
-
At first, let’s prove this fact: for any positive integer x and non-negative integer K it is always possible to divide numbers x^K, (x + 1)^K, (x + 2)^K, … , (x + 2^{(K + 1)} - 1)^K in two groups with equal sum of elements. Consider two groups of these numbers: A = {(x+i)^K | i has even number of bits in binary representation} and B = {(x+i)^K | i has odd number of bits in binary representation}.
-
For example, if K = 2 and x = 5 then A = {(5 + 0)^2, (5 + 3)^2, (5 + 5)^2, (5 + 6)^2} and B = {(5 + 1)^2, (5 + 2)^2, (5 + 4)^2, (5 + 7)^2}, as numbers 0, 3, 5, 6 have even number of bits, and numbers 1, 2, 4, 7 have odd number of bits. One may ensure, that indeed Sum(A) = 5^2 + 8^2 + 10^2 + 11^2 = 310 and Sum(B) = 6^2 + 7^2 + 9^2 + 12^2 = 310 as well.
-
Proof
Let’s prove that fact by mathematical induction on K. Case K = 0 is trivial. For K = 1 we will have A = {x, x + 3} and B = {x + 1, x + 2}, so Sum(A) = Sum(B) = 2x + 3. Now assume that we already prove our statement for K = 0, K = 1, …, K = M - 1 and we now want to prove statement for K = M. Notice, that for any i from [0; 2^{M}) we always have bits(i + 2^M) = bits(i) + 1. Let’s write down some equations:
Sum(B) - Sum(A) = - \sum_{i = x}^{x + 2^{M + 1} - 1} i^M * (-1)^{bits(i)} = \sum_{i = x}^{x + 2^M - 1} ((i + 2^M)^M - i^M) * (-1)^{bits(i)} = \sum_{i = x}^{x + 2^M - 1} (\sum_{j = 0}^{M - 1} i^j \binom{M}{j} (2^M)^{M - j})* (-1)^{bits(i)} = \sum_{j = 0}^{M - 1} \binom{M}{j} (2^M)^{M - j} \sum_{i = x}^{x + 2^M - 1} i^j * (-1)^{bits(i)} (here we use Binomial Theorem and then changed order of summation).
In the last equation, we will always have \sum_{i = x}^{x + 2^M - 1} i^j * (-1)^{bits(i)} = 0 for any
j from [0, M) by induction assumption, thus getting Sum(B) - Sum(A) = 0.Also, proof by brute-force is possible - there are only O(N) possible x for each K, thus we
can ensure that division is indeed correct by writing simple program. -
Now the question is: for how big values of N we can calculate dynamic programming, described previously? Let’s calculate it’s complexity. First dimension of dp will have N elements, while second should have 2 * (1^K + 2^K + ... + N^K), which is O(N^{K + 1}) (proof).
So in total, we will have O(N^{K + 2}) states and two transitions from each state (we can either add current number to the first group or to the second). After using bitsets, complexity will become O(N^{K+2} / 64) both time and memory. -
But as was said earlier, we need to calculate this dp only for “small” N. In turns out, that for N >= 2^{K + 1} answer will always be minimum possible (i.e. 1 for N = 1, 2 (modulo 4), and 0 otherwise).
-
Proof
Write described above dynamic programming for all 2^{K + 1} <= N < 2^{K + 2} and ensure that this is indeed true . If you have formal proof, please share it in comments!
-
So now the algorithm is as follows: while N >= 2^{K + 2} we will distribute 2^{K + 1} biggest elements in accordance to described strategy, and reduce N by 2^{K + 1}. At the end, we will have N < 2^{K + 2}, when we will run our dynamic programming part of solution. Answer restoration is really simple - just go backwards in dp. For better understanding, please refer to model implementation given below.
-
What about total complexity of solution? Obviously, first part is O(N) per testcase. Second part can be done just once as precalculation. But it’s complexity still is O(M^{K + 2} / 64), where M = 2^{K + 2}. This will be fast enough for K <= 3. However, for K = 4 we need will have M = 64, thus total complexity will be ~64^5, which is actually too much. But we can observe that we can actually reduce size of second dimension of dp - setting it’s size to some value 2 *S means that we are searching for such a solution, that all A[i] - B[i] are in range [-S; S]. You can experiment with this a little bit and ensure that actually S = 2*10^7 will be enough.
-
Also, as we actually need some small amount of values, we can precomputate them locally and hardcode them in solution .
ALTERNATE EXPLANATION:
SOLUTIONS:
Setter's Solution (C++)
#include <bits/stdc++.h>
using namespace std;
bool checker(int n, int k, int goal, string& s){
function<__int128(int)> power = [&](int x){
__int128 res = 1;
for(int i = 0; i < k; ++i){
res *= x;
}
return res;
};
__int128 is = 0;
for(int i = 0; i < n; ++i){
is += s[i] == '1' ? power(i + 1) : -power(i + 1);
}
return is == goal;
}
const int S = 20 * 1000 * 1000;
vector<bitset<S + S>> dp;
vector<int> answer, cost;
vector<vector<int>> belong;
void calculate_dp(int n, int k){
function<int(int)> power = [&](int x){
int res = 1;
for(int i = 0; i < k; ++i){
res *= x;
}
return res;
};
dp.resize(n);
answer.resize(n);
cost.resize(n);
belong.resize(n);
dp[0][S] = 1;
for(int i = 1; i < n; ++i){
cost[i] = power(i);
dp[i] = (dp[i - 1] << cost[i]) | (dp[i - 1] >> cost[i]);
}
for(int i = 1; i < n; ++i){
for(int j = S; ; ++j) if (dp[i][j]){
answer[i] = j - S;
vector<int>& res = belong[i];
for(int l = i; l >= 1; --l){
if (j - cost[l] >= 0 && dp[l - 1][j - cost[l]]){
res.push_back(l);
j -= cost[l];
}else{
assert(j + cost[l] < S + S);
assert(dp[l - 1][j + cost[l]]);
j += cost[l];
}
}
reverse(res.begin(), res.end());
assert(j == S);
break;
}
}
}
int main(){
int k; cin >> k;
calculate_dp(1 << (k + 2), k);
int tc; cin >> tc;
while(tc--){
int n; cin >> n;
string result(n, '0');
while(n >= 1 << (k + 2)){
for(int i = 0; i < 1 << (k + 1); ++i){
if (__builtin_popcount(i) & 1){
result[n - 1 - i] = '1';
}
}
n -= 1 << (k + 1);
}
for(auto i: belong[n]){
result[i - 1] = '1';
}
cout << answer[n] << endl;
cout << result << endl;
assert(checker(n, k, answer[n], result));
}
cerr << "Time elapsed: " << clock() / (double) CLOCKS_PER_SEC << endl;
return 0;
}
Tester's Solution (Kotlin)
fun Int.pow(exponent: Int): Int = this * (if(exponent == 1) 1 else pow(exponent - 1))
val precalced = listOf(
listOf(),
listOf(
Pair(0, ""),
Pair(1, "0"),
Pair(1, "10"),
Pair(0, "110"),
Pair(0, "0110"),
Pair(1, "11010"),
Pair(1, "111100"),
Pair(0, "0111100"),
Pair(0, "11011100"),
Pair(1, "111110100"),
Pair(1, "0111111000"),
Pair(0, "11011111000"),
Pair(0, "111110111000"),
Pair(1, "1111111110000"),
Pair(1, "11011111110000"),
Pair(0, "111110111110000"),
Pair(0, "1111111110110000"),
Pair(1, "11011111111010000"),
Pair(1, "111111011111010000"),
Pair(0, "0111111111100101000"),
Pair(0, "11111011111010101000"),
Pair(1, "111111111101010101000"),
Pair(1, "1110111111101010101000"),
Pair(0, "11111111011101010100100"),
Pair(0, "110111111110101010010100"),
Pair(1, "1111111011110101010010100"),
Pair(1, "10111111111010101001010100"),
Pair(0, "111111101111010010101010100"),
Pair(0, "1011111111101001010101010100"),
Pair(1, "11111101111101001010101010100"),
Pair(1, "011111111110100101010101010100"),
Pair(0, "1111110111101010101010101010100"),
Pair(0, "01111111110101010101010101010100"),
Pair(1, "111110111110101010101010101010100"),
Pair(1, "0111111111010101010101010101010010"),
Pair(0, "11111011111010101010101010101001010"),
Pair(0, "111111111101010101010101010101001010"),
Pair(1, "1111101111101010101010101010100101010"),
Pair(1, "11111111110101010101010101010100101010"),
Pair(0, "111101111110101010101010101010010101010"),
Pair(0, "1111111111010101010101010101001010101010"),
Pair(1, "11110111111010101010101010101001010101010"),
Pair(1, "111111111101010101010101010100101010101010"),
Pair(0, "1111011111101010101010101010010101010101010"),
Pair(0, "11111111110101010101010101001010101010101010"),
Pair(1, "111101111110101010101010101001010101010101010"),
Pair(1, "1111111111010101010101010100101010101010101010"),
Pair(0, "11110111111010101010101010010101010101010101010"),
Pair(0, "111111111011010101010101010010101010101010101010"),
Pair(1, "1111011111101010101010101001010101010101010101010"),
Pair(1, "11111111101101010101010101001010101010101010101010"),
Pair(0, "111011111110101010101010100101010101010101010101010"),
Pair(0, "1111111110110101010101010010101010101010101010101010"),
Pair(1, "11101111111010101010101010010101010101010101010101010"),
Pair(1, "111111111011010101010101001010101010101010101010101010"),
Pair(0, "1110111111101010101010100101010101010101010101010101010"),
Pair(0, "11111111101101010101010010101010101010101010101010101010"),
Pair(1, "111011111110101010101010010101010101010101010101010101010"),
Pair(1, "1111111110110101010101001010101010101010101010101010101010"),
Pair(0, "11101111111010101010100101010101010101010101010101010101010"),
Pair(0, "111111110111010101010100101010101010101010101010101010101010"),
Pair(1, "1110111111101010101010010101010101010101010101010101010101010"),
Pair(1, "11111111011101010101010010101010101010101010101010101010101010"),
Pair(0, "110111111110101010101001010101010101010101010101010101010101010"),
Pair(0, "1001100110011001100110011001100110011001100110011001100110011001")
),
listOf(
Pair(0, ""),
Pair(1, "0"),
Pair(3, "10"),
Pair(4, "110"),
Pair(2, "1110"),
Pair(3, "10110"),
Pair(1, "010110"),
Pair(0, "0010110"),
Pair(0, "10010110"),
Pair(1, "100111010"),
Pair(1, "1011011010"),
Pair(0, "01000111010"),
Pair(0, "111111110010"),
Pair(1, "1110111101100"),
Pair(1, "10011101011100"),
Pair(0, "001101111111000"),
Pair(0, "0101001111111000"),
Pair(1, "01100111011111000"),
Pair(1, "011011111111011000"),
Pair(0, "0011111111111110000"),
Pair(0, "11110011111111110000"),
Pair(1, "110111111101111110000"),
Pair(1, "1110111111111011110000"),
Pair(0, "11101011111111101110000"),
Pair(0, "101011111111111111100000"),
Pair(1, "0011111111101111111010000"),
Pair(1, "01100111111101111111100000"),
Pair(0, "001110111111111101111100000"),
Pair(0, "0111011111111111111101100000"),
Pair(1, "10011111011111111111110100000"),
Pair(1, "011101111111101111111110100000"),
Pair(0, "0011010111111111011111110100000"),
Pair(0, "11111110111111111111111001010000"),
Pair(1, "001101111111110111111110101010000"),
Pair(1, "1001011111111111111101110101001000"),
Pair(0, "01110111110111111111111010100101000"),
Pair(0, "101110111111111110111111010010101000"),
Pair(1, "1001111011111111111111100101010101000"),
Pair(1, "10110111111111101111111010101010101000"),
Pair(0, "010111111111111111111101010101010100100"),
Pair(0, "0001111111111101111111101010101010010100"),
Pair(1, "00111111111111111111101101010101001010100"),
Pair(1, "101110111110111111111110101010100101010100"),
Pair(0, "1011111110111111101111110101010010101010100"),
Pair(0, "10101011110111111111111010100101010101010100"),
Pair(1, "001111011111111111011111010010101010101010100"),
Pair(1, "0110111101111111111111101001010101010101010100"),
Pair(0, "00110101111111110111111101001010101010101010010"),
Pair(0, "000101111011111111111110100101010101010101001010"),
Pair(1, "1111110111111111101111110010101010101010101001010"),
Pair(1, "11011111011111111111111010010101010101010100101010"),
Pair(0, "110111111111111111101111001010101010101010010101010"),
Pair(0, "1111011011111111111111100101010101010101010010101010"),
Pair(1, "10111110111111110111111100101010101010101001010101010"),
Pair(1, "111111111101111111111110010101010101010100101010101010"),
Pair(0, "1001110111111111011111101010101010101010100101010101010"),
Pair(0, "01111110111111111111111001010101010101010010101010101010"),
Pair(1, "011011111111111111011110101010101010101001010101010101010"),
Pair(1, "1111011011111111111111100101010101010100101010101010101010"),
Pair(0, "00110111111111110111111010101010101010100101010101010101010"),
Pair(0, "111111111101111111111011010101010101010010101010101010101010"),
Pair(1, "0001111111111111101111101010101010101001010101010101010101010"),
Pair(1, "11101111110111111111101101010101010100101010101010101010101010"),
Pair(0, "011111011111110111111110101010101010100101010101010101010101010"),
Pair(0, "0110100101101001011010010110100101101001011010010110100101101001")
),
listOf(
Pair(0, ""),
Pair(1, "0"),
Pair(7, "10"),
Pair(18, "110"),
Pair(28, "1110"),
Pair(25, "11110"),
Pair(7, "101110"),
Pair(26, "1110001"),
Pair(4, "11001001"),
Pair(5, "100101001"),
Pair(1, "0111111010"),
Pair(12, "11110010110"),
Pair(0, "001011100110"),
Pair(1, "1111101111100"),
Pair(1, "10100110010110"),
Pair(0, "001111010111100"),
Pair(0, "0110100110010110"),
Pair(1, "00010110111110010"),
Pair(1, "101100101010100101"),
Pair(0, "1100110110100111100"),
Pair(0, "01011111110111111000"),
Pair(1, "100111111110101101100"),
Pair(1, "1111111111111110111000"),
Pair(0, "10001111101111110111000"),
Pair(0, "101011110110111110111000"),
Pair(1, "1111111110110111111011000"),
Pair(1, "10110011010111111111110000"),
Pair(0, "011101101111111011111110000"),
Pair(0, "0010101110110110111111110000"),
Pair(1, "00000010010111111110111110000"),
Pair(1, "001100100101111111111101110000"),
Pair(0, "1011101001011111011111011110000"),
Pair(0, "11100001111110111111111111100000"),
Pair(1, "100101010101111101111111111100000"),
Pair(1, "1101101011110111111101111111100000"),
Pair(0, "11001010110111111111111101111100000"),
Pair(0, "000011001101111111111111111011100000"),
Pair(1, "0001011111101111111111111111110100000"),
Pair(1, "10000011010111110111111111111110100000"),
Pair(0, "111011111111111111111110111111110100000"),
Pair(0, "0101111110110111111111111110111110010000"),
Pair(1, "10110011111010111111111111111110101010000"),
Pair(1, "111100101100111111101111111111110101010000"),
Pair(0, "0110110101110111101111110111111110101001000"),
Pair(0, "10101111111111111101111111111111010010101000"),
Pair(1, "101100110101111111111111110111110101010101000"),
Pair(1, "1010011111101111111111111111110110101010101000"),
Pair(0, "00001111100111111111111101111111010101010010100"),
Pair(0, "010101111111101111111111111111011010101001010100"),
Pair(1, "0111111011111111111111011111111101010100101010100"),
Pair(1, "10111110111011111111111111110111101010010101010100"),
Pair(0, "010010110111111101111111111111110101001010101010100"),
Pair(0, "0000110010011110111111110111111110010101010101010100"),
Pair(1, "10110100011111101111111111111111010010101010101010010"),
Pair(1, "001010011111111111111111111011111010010101010101001010"),
Pair(0, "1011101011001111011111111111111101001010101010100101010"),
Pair(0, "01111111010111111111111111011111100101010101010100101010"),
Pair(1, "101100101100111111111111111111110100101010101010010101010"),
Pair(1, "1111100111011111111111111110111110010101010101001010101010"),
Pair(0, "10010111111101111111111111111111001010101010101001010101010"),
Pair(0, "010011001111110111111111101111110101010101010100101010101010"),
Pair(1, "0010011001110111111111111111111100101010101010010101010101010"),
Pair(1, "11010110111010111111111110111111010101010101001010101010101010"),
Pair(0, "000011101111110111111111111111110010101010100101010101010101010"),
Pair(0, "1001011001101001100101100110100110010110011010011001011001101001")
),
listOf(
Pair(0, ""),
Pair(1, "0"),
Pair(15, "10"),
Pair(64, "110"),
Pair(158, "1110"),
Pair(271, "11110"),
Pair(317, "111110"),
Pair(126, "1111110"),
Pair(34, "11010001"),
Pair(253, "111110001"),
Pair(13, "1010110110"),
Pair(92, "11111010110"),
Pair(30, "110010101001"),
Pair(47, "1101001010110"),
Pair(31, "11111110000101"),
Pair(2, "100001110101001"),
Pair(0, "0000111000001110"),
Pair(1, "10000111111101010"),
Pair(13, "111111001011110001"),
Pair(0, "1000010010110111100"),
Pair(0, "00001010110111011100"),
Pair(9, "111111111011011011100"),
Pair(1, "0101001111011111110100"),
Pair(0, "00001000100110111101100"),
Pair(0, "000010011110110111110100"),
Pair(1, "1000110101011101111111000"),
Pair(5, "11101011111011101111001010"),
Pair(0, "001111000000011110111111000"),
Pair(0, "0001011111000101111011111000"),
Pair(5, "11101111111011001110111010001"),
Pair(1, "100010010001010111111111011000"),
Pair(0, "1101100001110111011111111101000"),
Pair(0, "00010000010110101111111111110000"),
Pair(1, "110101011100010111011111111110000"),
Pair(1, "1010101110111111111011101111101000"),
Pair(0, "01000111001011010111111101111110000"),
Pair(0, "100101010100011101111111111011110000"),
Pair(1, "1111111010111011111110111111111000100"),
Pair(1, "10010100100101011111110111111101110000"),
Pair(0, "000000100100010111110111111111111010000"),
Pair(0, "0000001001001100010111111111111111100000"),
Pair(1, "01000101010110011111111111011111111100000"),
Pair(1, "011111101111111010111110101111111111100000"),
Pair(0, "0001010011000001110111111111111011111100000"),
Pair(0, "00000100000101011110111111111111110111100000"),
Pair(1, "001111101111101111111010111111111110111100000"),
Pair(1, "0011010001010111011111111111011111111101010000"),
Pair(0, "11000100011101010101110111110111111111101010000"),
Pair(0, "010100000101011110011111111111111111011101001000"),
Pair(1, "0111110100001101010111111101111111111110010101000"),
Pair(1, "11010111011111111011111111111111110111101010101000"),
Pair(0, "010111100101110101110111111111111111110101010100100"),
Pair(0, "0001010001000011100101111111011111111110101010010100"),
Pair(1, "10011111101101111111101111111111111111010100101010100"),
Pair(1, "010001110010110101110111111101111111111010010101010100"),
Pair(0, "0101000001010100000101001111110111111110101010101010100"),
Pair(0, "11010001000110010111111111111011111111100101010101010010"),
Pair(1, "000100010000011001110101111111111101111100101010101001010"),
Pair(1, "1111101010101110111111111111111011111100110101010101001010"),
Pair(0, "01010001010001010101101111111111110111101010101010100101010"),
Pair(0, "011000010001010001111111110111111111111001010101010010101010"),
Pair(1, "1011101100101111101011111111111111110110101010101001010101010"),
Pair(1, "00000001000001011001010111111111111111100101010100101010101010"),
Pair(0, "010010000000010101011111111111111011111010101010100101010101010"),
Pair(0, "0110100110010110100101100110100101101001100101101001011001101001")
)
)
class EQDIV_SOLVER(val K: Int) {
val tb = precalced[K]
init {
require(tb.lastIndex == 64)
}
fun solve(N: Int): Pair<Int, String> {
if (N <= 32) {
return tb[N]
} else {
val (diff, partition) = tb[32 + N % 32]
return Pair(diff, partition + tb[64].second.take(32).repeat((N - 32) / 32))
}
}
}
fun main(args: Array<String>) {
val br = java.io.BufferedReader(java.io.InputStreamReader(System.`in`))
val bw = java.io.BufferedWriter(java.io.OutputStreamWriter(System.`out`))
val K = br.readLine()!!.toInt()
require(K in 1..4)
val T = br.readLine()!!.toInt()
require(T in 1..2000)
val solver = EQDIV_SOLVER(K)
var sumN = 0L
repeat(T) {
val N = br.readLine()!!.toInt()
sumN += N
val (diff, ans) = solver.solve(N)
bw.write("$diff\n$ans\n")
}
require(sumN <= 5000000L)
bw.flush()
}