【算法】从 n 个数里面随机取 m 个数

Table of Contents

题目简介
#

最近面试 Two Sigma，电话面试问了一个很有趣的问题：给一个函数 getRandom(x) 能随机返回[0,x) 中的一个数，求实现一个 RandomNumberGenerator(n) 数据结构，里面有一个 generate() 接口，要求每次调用随机返回[0,n) 中的一个数，同时不能返回已经返回过的数。

题目分析
#

这个题换一种说法其实就是从 n 个数里面随机取 m 个数。工作中经常遇到类似的问题，一般调用已有库函数就可以实现，所以从来没有想过这个方法的具体实现。

最直接解法
#

用一个 Set 来保存已经出现的数，generate() 函数不停地调用 getRandom(n)，直到生成的数字不在 set 里面，返回这个数，同时把这个数加进 Set

public class RandomNumberGenerator {
private final int n;
private final Set<Integer> blacklist;

	public int generate() {
		if (blacklist.size() == n) {
			return -1;
		}
		while (true) {
			int next = getRandom(n);
			if (!blacklist.contains(next)) {
				blacklist.add(next);
				return next;
			}
		}
	} 
}

这个解法最坏情况可能导致无限死循环，假设 blacklist 里面数字个数为 m，那么每次成功的概率为 (n - m) / n，如果 n 和 m 十分接近，则这个概率会很小。

改进版直接解法
#

上面方法的问题在于过于随机性，我们需要一个运行时间是确定的解法。可以稍微调整一下思路，我们可以用一个数组 list 来保存所有没有被访问过的数，generate() 函数只需要调用 getRandom(list.size()) 一次，得到返回数的 index，通过 index 将这个数从数组移除，同时返回这个数。

public class RandomNumberGenerator {
private final List<Integer> numbers;

	public int generate() {
		if (numbers.size() == 0) {
			return -1;
		}
		int index = getRandom(numbers.size());
		return numbers.remove(index);
	} 
}

时间复杂度是 O(n)，因为每次 remove 后需要调整。空间复杂度是 O(n)。如果输入的 n 很大浪费时间空间。

稍微优化算法
#

可以先考虑优化时间复杂度。可以看到上一个解法时间复杂度之所以是 O(n)，主要是由于 remove 后需要整体调整导致的，那有什么办法可以减少这里的操作呢？其实我们不需要 remove 后整体往前移动，只需要 swap 交换被选中的 index 元素和 List 最后没有被交换的元素即可，这样交换只需要 O(1)

public class RandomNumberGenerator {
private final List<Integer> numbers;
// init cur = n
private int cur;

	public int generate() {
		if (cur == 0) {
			return -1;
		}
		int index = getRandom(cur);
		int reuslt = numbers.get(index);
		// swap index number with the last number
		numbers.set(index, numbers.get(cur - 1));
		numbers.set(cur - 1, result);
		cur--;
		return result;
	} 
}

这个算法时间复杂度缩减为 O(1)，但是空间复杂度还是 O(n)。

优化算法
#

通常从 n 个数里面随机取 m 个数，m 都远小于 n，所以对于上面的解法我们还可以对空间复杂度进行优化。

我们如果研究下上面的 numbers 数组，可以发现如果一个 index 从来没有被选过，那么他所对应的数与 index 想同，所以我们其实只需要一个数据结构来存 index 对应值与 index 不相同的 index 与其值（很自然相当 key-value）

public class RandomNumberGenerator {
private final Map<Integer, Integer> indexToNumber;
// init cur = n
private int cur;

	public int generate() {
		if (cur == 0) {
			return -1;
		}
		int index = getRandom(cur);
		int reuslt = indexToNumber.getOrDefault(index, index);
		numbers.set(index, indexToNumber.getOrDefault(cur - 1, cur - 1));
		cur--;
		return result;
	} 
}

算法时间复杂度是 O(1)，空间复杂度是 O(m)

题目简介 #

题目分析 #

最直接解法 #

改进版直接解法 #

稍微优化算法 #

优化算法 #