【LeetCode】430. Rand7 实现 Rand10

优质

小牛编辑

131浏览

2023-12-01

更多面试题总结请看：【面试题】技术面试题汇总

引言

概率生成问题是一种比较常见的面试题，常见题型举例：

有一枚不均匀的硬币，要求产生均匀的概率分布
有一枚均匀的硬币，要求产生不均匀的概率分布，如 0.25 和 0.75
利用 Rand7() 实现 Rand10()

本文将依次讲解这三个问题，然后总结此类问题的通用方法。

不均匀硬币，产生等概率

现有一枚不均匀的硬币 coin()，能够返回 0、1 两个值，其概率分别为 0.6、0.4。要求使用这枚硬币，产生均匀的概率分布。即编写一个函数 coin_new() 使得它返回 0、1 的概率均为 0.5。

// 不均匀硬币，返回 0、1 的概率分别为 0.6、0.4
int coin() {
    return (rand() % 10) > 5;
}

统计抛两次硬币的结果的概率分布：

结果	0	1
0	0.6*0.6=0.36	0.6*0.4=0.24
1	0.4*0.6=0.24	0.4*0.4=0.16

不难发现，连续抛两枚硬币得到 0 1 和 1 0 的概率分布是相同的。因此这道题的解法就是连续抛两次硬币，如果得到 0 1，返回 0；如果得到 1 0，返回 1；如果两次结果相同，则重新抛。

以此类推，无论这枚不均匀硬币的概率是多少，都可以用这种方法得到等概率的结果。

代码如下：

int coin_new() {
    while (true) {
        int a = coin();
        if (coin() != a) return a;
    }
}

完整测试代码：

#include <iostream>
#include <map>
using namespace std;

// 不均匀硬币
int coin() {
    return (rand() % 10) > 4;
}

int coin_new() {
    while (true) {
        int a = coin();
        if (coin() != a) return a;
    }
}

// 测试
int main() {
    int N = 1000000;
    float count[2] = {0, 0};
    for (int i = 0; i < N; i++)
        count[coin_new()] ++;

    cout << "0: " << count[0]/N << endl;
    cout << "1: " << count[1]/N << endl;
}

输出：

0: 0.50037
1: 0.49963

均匀硬币，产生不等概率

现有一枚均匀的硬币 coin()，能够返回 0、1 两个值，其概率均为 0.5。要求编写一个函数 coin_new()，使得它返回指定的 0、1 概率分布。

// 均匀硬币
int coin() {
    return rand() % 2;
}

P(0) = 1/4，P(1) = 3/4

对于均匀硬币而言，连续抛两次，得到 0 0、0 1、1 0、1 1 的概率均为 1/4。显然，只需要连续抛两次硬币，如果得到 0 0，返回 0，其他情况返回 1。

int coin_new() {
    return coin() || coin();
}

测试输出：

0: 0.249249
1: 0.750751

P(0) = 1/3，P(1) = 2/3

连续抛两次硬币。如果得到 1 1，返回 0；如果得到 1 0 或 0 1，返回 1；如果得到 0 0，继续抛硬币。

int coin_new() {
    while (true) {
        int a = coin(), b = coin();
        if (a && b) return 0;
        if (a || b) return 1;
    }
}

测试输出：

0: 0.333663
1: 0.666337

P(0) = 0.3，P(1) = 0.7

每抛一次硬币，会得到二进制数的一位
连续抛 4 次硬币，可以等概率生成 [0, 15] 的每个数，记为 $x$
去掉 [10, 15]，剩下 [0, 9] 的每个数依然是等概率的
如果 $x \in [0, 2]$，返回 0；$x \in [4, 9]$，返回 1；$x ≥ 10$，重复上述过程

int coin_new() {
    while (true) {
        int x = 0;
        for (int i = 0; i < 4; i++) {
            x = (x << 1) + coin();
        }
        if (x <= 2) return 0;
        if (x <= 9) return 1;
    }
}

测试输出：

0: 0.300324
1: 0.699676

总结

不难发现，上面这三道题目其实都是同一种思路：

每抛一次硬币，会得到二进制数的一位
连续抛 $k$ 次硬币，可以等概率生成 $[0, 2^k-1]$ 的每个数
在 $[0, 2^k-1]$ 中，选取 $m$ 个数返回 0，$n$ 个数返回 1，则 0、1 的概率分别为 $\frac{m}{m+n}$、$\frac{n}{m+n}$

关于 $k$ 的选择，最少需要满足 $ N <= 2^k-1$，N 是生成对应概率分布至少需要多少个不同数字。比如要生成 1/3、2/3 的分布，至少需要 3 个不同数字，则 $N = 3, k = 2$；要生成 3/10、7/10 的分布，至少需要 10 个数字，则 $N = 10, k = 4$。

$k$ 最多则没有限制，我们总可以通过抛更多次硬币来解决问题，只需要把无用的数字舍弃即可。但我们的目的是尽可能减少无用数字的比例，因为每次遇到无用数字时，都需要重新生成新的数字。

Rand7 生成 Rand10

这道题是 LeetCode 430 题。已有方法 Rand7() 可生成 1 到 7 范围内的均匀随机整数，试写一个方法 Rand10() 生成 1 到 10 范围内的均匀随机整数。

我们先从抛硬币开始入手。抛硬币可以看作是 Rand2()，均匀生成 0、1 两个整数。那么，如何根据 Rand2() 生成 Rand10()？按照前文的描述，我们只需要将每次抛硬币的结果，看作二进制的每一位，就可以得到 $[0, 2^k-1]$ 范围内的均匀随机整数。对于这道题，只需要抛 4 次硬币，就能得到 [0, 15] 范围的整数，这包含了我们需要的范围 [1, 10]。因此只需要在得到 [1, 10] 范围的整数时返回，其他情况则重新抛硬币。

int rand10() {
    while (true) {
        int x = 0;
        for (int i = 0; i <4; i++) x = x << 1 + rand2();
        if (1 <= x <= 10) return x;
    }
}

再回到这道题，我们可以将 Rand7() 视作一个 7 进制位的生成器。由于 Rand7() 的结果范围是 [1, 7]，故我们取 Rand7() - 1 作为对应的 7 进制位。每执行 $k$ 次 Rand7()，将得到一个 $k$ 位的 7 进制整数，在 $[0, 7^k-1]$ 范围内均匀分布。

对于本题而言，只需执行 $k = 2$ 次 Rand7()，就可以得到范围为 $[0, 48]$ 的均匀整数：

当 $x \in [1, 10]$ 时返回 $x$，否则重新计算：

int rand10() {
    while (true) {
        int x = (rand7() - 1) * 7 + (rand7() - 1);
        if (x >= 1 && x <= 10) return x;
    }
}

以上代码运行时间为 24ms。

进一步优化

上面这种方法效率很低，因为我们只需要 [1, 10] 范围的整数，但会生成 [0, 48] 范围的整数，相当于有 4/5 的数字都是无用的，需要重新生成。

可不可以充分利用 [0, 48] 范围内尽可能多的数字呢？这样可以让 while 舍弃的数字比例更少，也就是让 while 重新执行的概率更低。

一种方法是选择 [1, 40] 范围里的数，通过取余运算来得到 [1, 10] 范围的数：

int rand10() {
    while (true) {
        int x = (rand7() - 1) * 7 + (rand7() - 1);
        if (x >= 1 && x <= 40)
            return x % 10 + 1;
    }
}

这种方法只有 0、41、42、…、48 一共 9 个无用数字，运行时间从 24ms 降低至 16ms。

我们还可以进一步优化。对于上面这 9 个无用数字，计算 x % 40 可以得到 [0, 8] 范围的均匀随机整数。此时再调用一次 Rand7()，计算 (x % 40) * 7 + Rand7()，这相当于 Rand8() * 7 + Rand7()。显然，我们可以得到 [1, 63] 范围的均匀随机整数。这时 [1, 60] 范围里的数都可以用来作取余运算，只有 61、62、63 共 3 个无用数字：

while (true) {
    int x = (rand7() - 1) * 7 + (rand7() - 1); // 0~48
    if (x >= 1 && x <= 40) return x % 10 + 1;

    x = (x % 40) * 7 + rand7(); // 1~63
    if (x <= 60) return x % 10 + 1;
}

这还不算完。对于 61、62、63，我们也可以用相同的思路，再调用一次 Rand7()，计算 (x - 61) * 7 + Rand7()，相当于 Rand2() * 7 + Rand7()，可以得到 [1, 21] 范围的均匀随机整数，这时再作取余运算，只有 1 个无用数字（21）：

int rand10() {
    while (true) {
        int x = (rand7() - 1) * 7 + (rand7() - 1); // 0~48
        if (x >= 1 && x <= 40) return x % 10 + 1;

        x = (x % 40) * 7 + rand7(); // 1~63
        if (x <= 60) return x % 10 + 1;

        x = (x - 61) * 7 + 7; // 1~21
        if (x <= 20) return x % 10 + 1;
    }
}

运行时间从 16ms 降低至 8ms。每次 while 执行的时候，只有 1 个无用数字（21）会被舍弃，重新执行的概率很低。

最后，注意这里每次乘的系数都是 7，这是为了保证等概率。乘的系数大于 7 或者小于 7，都无法保证等概率：

((rand7() - 1) * 6) + rand7() - 1：6、12、18… 概率偏高
((rand7() - 1) * 8) + rand7() - 1：永远无法生成 7 的倍数

RandM 生成 RandN

已知 RandM() 可以等概率的生成 [0, M-1] 范围的随机整数，那么执行 k 次，每次都得到 M 进制的一位，可以等概率生成 $[0, M^k-1]$ 范围的随机整数，记为 $x$
RandN 至少需要 N 个均匀随机整数，因此只需要取 $k$，使得 $M^k-1 >= N$ 即可
此时有多种方式得到 RandN：
- 一种是只在 $x \in [0, N-1]$ 时返回 $x$
- 另一种是利用取余运算，在保证等概率的前提下，尽可能多的利用生成的数字，从而减少舍弃的数字比例，降低 while 重复执行的概率

附录：测试代码

均匀硬币产生不等概率：

#include <iostream>
#include <map>
using namespace std;

int coin() {
    return rand() % 2
}

// TODO：替换为具体实现
int coin_new() {}

// 测试
int main() {
    int N = 1000000;
    float count[2] = {0, 0};
    for (int i = 0; i < N; i++)
        count[coin_new()] ++;

    cout << "0: " << count[0]/N << endl;
    cout << "1: " << count[1]/N << endl;
}