Number of k-good subarrays
Contest: Codeforces Round 955 (Div. 2, with prizes from NEAR!)
Difficulty rating: 2300
Tags: bitmasks, brute force, combinatorics, divide and conquer, dp, math, meet-in-the-middle

Let $$$bit(x)$$$ denote the number of ones in the binary representation of a non-negative integer $$$x$$$.

A subarray of an array is called $$$k$$$-good if it consists only of numbers with no more than $$$k$$$ ones in their binary representation, i.e., a subarray $$$(l, r)$$$ of array $$$a$$$ is good if for any $$$i$$$ such that $$$l \le i \le r$$$ condition $$$bit(a_{i}) \le k$$$ is satisfied.

You are given an array $$$a$$$ of length $$$n$$$, consisting of consecutive non-negative integers starting from $$$0$$$, i.e., $$$a_{i} = i$$$ for $$$0 \le i \le n - 1$$$ (in $$$0$$$-based indexing). You need to count the number of $$$k$$$-good subarrays in this array.

As the answer can be very large, output it modulo $$$10^{9} + 7$$$.

Input format:
Each test consists of multiple test cases. The first line contains an integer $$$t$$$ ($$$1 \le t \le 10^{4}$$$) — the number of test cases. The following lines describe the test cases.

The single line of each test case contains two integers $$$n$$$, $$$k$$$ ($$$1 \le n \le 10^{18}, 1 \le k \le 60$$$).

Output format:
For each test case, output a single integer — the number of $$$k$$$-good subarrays modulo $$$10^{9} + 7$$$.

Constraints:
Time limit: 2.0 s
Memory limit: 256.0 MB

Note:
For the first test case $$$a = [0, 1, 2, 3, 4, 5]$$$, $$$k = 1$$$.

To find the answer, let's write all the numbers in binary representation:

$$$$$$a = [\color{green}{000}, \color{green}{001}, \color{green}{010}, \color{red}{011}, \color{green}{100}, \color{red}{101}]$$$$$$

From this, it can be seen that the numbers $$$3$$$ and $$$5$$$ have $$$2 \ge (k = 1)$$$ ones in their binary representation, so the answer should include all subarrays that do not contain either $$$3$$$ or $$$5$$$, which are the subarrays (in $$$0$$$-based indexing): ($$$0$$$, $$$0$$$), ($$$0$$$, $$$1$$$), ($$$0$$$, $$$2$$$), ($$$1$$$, $$$1$$$), ($$$1$$$, $$$2$$$), ($$$2$$$, $$$2$$$), ($$$4$$$, $$$4$$$).

Example input:
10
6 1
16 2
1 1
3 1
31 3
14 1
1337 5
100000 20
795569939321040850 56
576460752303423268 59

Example output:
7
35
1
6
155
8
7323
49965
741136395
66679884

Editorial:
Let $$$f(n, k)$$$ be a function to solve the problem, which will return three values $$$(l, r, ans)$$$ such that:
• $$$l$$$ — for the first $$$l$$$ numbers (i.e. from $$$0$$$ to $$$l - 1$$$) $$$bit(x) \le k$$$ holds, while $$$bit(l) > k$$$ or $$$l \ge n$$$;
• $$$r$$$ — for the last $$$r$$$ numbers (i.e. from $$$n - r$$$ to $$$n - 1$$$) $$$bit(x) \le k$$$ holds, while $$$bit(n - r - 1) > k$$$ or $$$n - r - 1 < 0$$$;
• $$$ans$$$ — the answer for the pair $$$(n, k)$$$.
Let's use the "divide and conquer" approach. Divide the segment from $$$0$$$ to $$$n$$$ into two segments:
1. from $$$0$$$ to $$$2^{m} - 1$$$, where $$$m$$$ is the maximum such that $$$2^{m} < n$$$;
2. from $$$2^{m}$$$ to $$$n - 1$$$.
It can be noticed that instead of calculating the function on the segment from $$$2^{k}$$$ to $$$n - 1$$$, it is possible to calculate it from $$$0$$$ to $$$n - 1 - 2^{m}$$$, but with a smaller $$$k$$$.
More formally, $$$f(n, k)$$$ can be recalculated through $$$f(2^{m}, k)$$$ and $$$f(n - 2^{m}, k - 1)$$$.
Based on the pair of triples $$$(l_{1}, r_{1}, ans_{1})$$$, which is obtained after calculating $$$f(2^{m}, k)$$$ and $$$(l_{2}, r_{2}, ans_{2})$$$, which is obtained after calculating $$$f(n - 2^{m}, k - 1)$$$, it is quite trivial to recalculate the triple for $$$f(n, k)$$$ using the formula in $$$O(1)$$$.
However, the solution still works in $$$O(n)$$$ time, as we will have a complete binary tree of recursive calls.
To fix this, it is only necessary to precalculate all $$$f(2^{m}, k)$$$ for each possible $$$k$$$ and $$$m$$$. For $$$f(2^{m}, k)$$$, it can be noticed that it is recalculated through $$$f(2^{m - 1}, k)$$$ and $$$f(2^{m - 1}, k - 1)$$$, so all values can be calculated using dynamic programming.
After this, in the recursive function, we always calculate $$$f(2^{m}, k)$$$ in $$$O(1)$$$, and for the next call, the first argument is halved, so now everything works in $$$O(\log n)$$$ for one test case.
In total, the solution is in $$$O(\log n)$$$ for the test case.