【数位DP】二进制问题

二进制问题 - 蓝桥云课

---

题目描述

小蓝最近在学习二进制。他想知道 $1$ 到 $N$ 中有多少个数满足其二进制表示中恰好有 $K$ 个 $1$ 。你能帮助他吗？

输入描述

输入一行包含两个整数 $N$ 和 $K$ 。

输出描述

输出一个整数表示答案。

输入输出样例

示例

输入

7 2

输出

3

评测用例规模与约定

对于 $30$ % 的评测用例， $1 ≤ N ≤ 10^6, 1 ≤ K ≤ 10$ 。

对于 $60$ % 的评测用例， $1 ≤ N ≤ 2 × 10^9, 1 ≤ K ≤ 30$ 。

对于所有评测用例， $1 ≤ N ≤ 10^{18}, 1 ≤ K ≤ 50$ 。

运行限制

• 最大运行时间：1s
• 最大运行内存: 256M

解题

方法一：数位DP

思路

求 $1$ 到 $N$ 中满足某种条件的数的数量，这是经典的数位DP问题。

对于本题来说，把 $N$ 转为二进制并计 $N$ 的二进制形式长度为 $len$ ，从最高位向低位开始递归枚举，并记录：

• 当前枚举的是第几位（idx）。
• 当前二进制表示中存在多少个 $1$ （cnt）。
• 当前这一位可以填的上限是多少。
  • 拿十进制来举例：如果右边界为 $987654$ ，现在枚举到了 $987x??$ 此时 $x$ 能填的数只有 $[0, 6]$ ，否则就会超过边界。二进制也是相同，假如右边界是 $110100$ ，现在枚举到了 $11x???$ ，此时 $x$ 能填的数只有 $0$ ，填 $1$ 就超过边界了。
  • 这一以来就明确了：需要一个变量（lim）记录之前的所有数字是否都达到了上限。

根据上面的思路写出 DFS ：

static long dfs(int idx, int cnt, boolean lim) {
    if (idx < 0) return cnt == k ? 1 : 0; // 递归出口  [0...len]所有位已经填完了时
    int up = lim ? digits[idx] : 1; // 上界  如果之前的数位都达到了上限 这位数最大就只能填{N在这一位的数}
    long res = 0L; // 方案数
    for (int i = 0; i <= up; ++i) { // 枚举
        // 数位向下一位  如果当前选的数是1就把计数增加  如果之前达到了上限且这一位也是上限就限制下一位的上限
        res += dfs(idx - 1, cnt + (i == 1 ? 1 : 0), lim && i == up);
    }
    return res;
}

发现在递归的过程中会出现很多重复的情况（特别是在 lim == false 时），所以使用记忆化递归优化（具体来说，增加一个 f 数组，第一维为 idx 第二维为 cnt ，值用于记录方案数即可）。

代码

import java.util.*;
import java.io.*;

public class Main {
    static final int N = 70;
    static long n;
    static int k;
    static int[] digits = new int[N];
    static long[][] f = new long[N][N];

    static long dfs(int idx, int cnt, boolean lim) {
        if (idx < 0) return cnt == k ? 1 : 0;
        if (!lim && f[idx][cnt] != -1) return f[idx][cnt];
        int up = lim ? digits[idx] : 1;
        long res = 0L;
        for (int i = 0; i <= up; ++i) {
            res += dfs(idx - 1, cnt + (i == 1 ? 1 : 0), lim && i == up);
        }
        return f[idx][cnt] = res;
    }

    static long count(long n) {
        for (int i = 0; i < N; ++i) Arrays.fill(f[i], -1L);
        int idx = 0;
        while (n > 0) {
            digits[idx++] = (int) (n & 1);
            n >>= 1;
        }
        return dfs(idx, 0, true);
    }

    public static void main(String[] args) throws IOException {
        Scanner in = new Scanner(System.in);
        n = in.nextLong();
        k = in.nextInt();
        System.out.println(count(n));
    }
}