PROBLEM LINK:
Author: Piyush Kumar
Tester: Minako Kojima
Editorialist: Pawel Kacprzak
DIFFICULTY:
EASY-MEDIUM
PREREQUISITES:
Bits, Recursion, DP, Memoization
PROBLEM:
You are given two integers N and K, where n ≤ 1018 and k = 3 or 4. The task is to count number of k-tuples a1, a2, …, ak where a1 + a2 + … + ak = n
and for each i = 0, …, k - 1, ai ^ ai+1 = ai+1. We call the first condition the pairwise condition and the second the sum condition.
QUICK EXPLANATION:
In order to fulfil the pairwise condition, you can notice that there are only 4 (k = 3) or 5 (k = 4) valid combinations for every bit of a tuple elements.
Based on that fact, you can use recursion starting from the leftmost bit of n and try to cover 1’s with possible valid combinations of corresponding bits of elements a tuple. If you decide to not cover a particular 1 with a weight b, you carry it to the next position and try to cover 2 ones with a weight b - 1 instead.
EXPLANATION:
For k = 3 there exists a very simple solution (see Alternative solution), but the below method is general and works for any k, even bigger than 4.
Without a loss of generality we assume that k = 3. The solution for k = 4 is analogous.
We say that a tuple is valid if the pairwise condition is fulfilled.
Let’s consider the i-th bit. There are only 4 combinations of bits at position i in a valid tuple:
1 2 3 4
a_1: 0, 1, 1, 1
a_2: 0, 0, 1, 1
a_3: 0, 0, 0, 1
You can see that any tuple formed by values in above columns is a valid tuple and moreover any other tuple is invalid.
In order to fulfil the sum condition, i.e. a1 + a2 + a3 = n, we can use recursion starting from the leftmost bit of n and try to cover 1’s in the binary representation of n with the above combinations. Any 1 at position i can be covered by elements of a tuple at position i or later. Let’s define a recursive function:
solve(b, carry) := number of valid tuples covering both: the rightmost b bits of n and a value carry * 2^b.
Let’s see an example for n = 3:
bit 2 1 0 n = 2 1 0 1
We start with b = 2 i.e. from the leftmost bit and call solve(2, 0), because there is no carry at the beginning. We see there is a 1here, so we have to cover it in a tuple. In order to do that, we can use the second combination from the above table and add solve(1, 0) to the result or decide to cover that bit later and add solve(1, 2) to the result covering a bit with a weight 2 with 2 bits with weight 1. Let’s consider the second situation. There is a 0 at position 1, so we have to cover only carried bits. We can use the third combination from the table to cover all of them or the second one to cover only one or the first one to decide to cover them later. Depending on our choice, we call the next recursion.
The bases case of recursion is when there are no more bits to cover, so when b = -1. In this situation, we return 1 if carry = 0 otherwise we return 0
In my solution, maximal carry is 4 * 64 in a case where k = 4, because n does not have more than 64 bits and we have exactly 4 elements in a tuple, so there are no more than 4 * 64 positions where we can put a 1, but with a more precise analysis, you can use a smaller maximal carry.
In order to avoid computing the result form solve(b, carry) many times, we use memoization to store already computed results.
Pseudo code
solve(b, carry)
if(b == -1)
return (carry == 0)
if(carry > max_carry)
return 0
if(mem[b][carry] is already computed)
return mem[b][carry]
new_carry = carry
if(n & (1LL << b))
new_carry++
for(d = 0; d <= min(carry, k); ++d)
mem[b][carry] += solve(b - 1, 2 * (new_carry - d))
return mem[b][carry]
Time complexity
O(64 * 64 * k) for one test case.
Alternative Solution:
For k = 3 there exists a simple formula that the result is 1 + (n/2). You can see a tester solution for more details.
AUTHOR’S AND TESTER’S SOLUTIONS:
Author’s solution can be found here.
Tester’s solution can be found here.