写在前面

面试的时候可能会这样的算法题：从无限的字符流中, 随机选出 10 个字符

相同的问题还有：

1.给定一个数据流，数据流长度N很大，且N直到处理完所有数据之前都不可知，
请问如何在只遍历一遍数据（O(N)）的情况下，能够随机选取出m个不重复的数据

2.在不知道文件行数的情况下，如何在只遍历一遍文件的情况下，随机选取出m行

分析

这个问题其实是问的是从n个样本中随机抽取m个，如何确保每个样本抽到的概率相同。这个抽样跟普通抽样不一样的是，不知道样本数据n有多大，或者说无限大（大到内存一次性放不下），这种情况改怎么确保随机性和每个样本数据被抽到的概率相同。

这是经典的蓄水池样本抽要算法，算法步骤如下：

如果接收的数据量小于m，则依次放入蓄水池。
当接收到第i个数据时，i >= m，在[0, i]范围内取以随机数d，若d的落在[0, m-1]范围内，则用接收到的第i个数据替换蓄水池中的第d个数据。
重复步骤2。

算法的精妙之处在于：当处理完所有的数据时，蓄水池中的每个数据都是以m/N的概率获得的。 证明算法自己查。

代码

/*
蓄水池抽样算法

场景：
1.给定一个数据流，数据流长度N很大，且N直到处理完所有数据之前都不可知，
请问如何在只遍历一遍数据（O(N)）的情况下，能够随机选取出m个不重复的数据
2.在不知道文件行数的情况下，如何在只遍历一遍文件的情况下，随机选取出m行
*/

type Node struct {
	val  int
	next *Node
}

func Rand(head *Node, n int) []int {
	result := make([]int, 0, n)
	idx := 0
	p := head
	rand.Seed(time.Now().Unix())
	for p != nil {
		if idx < n {
			idx += 1
			result = append(result, p.val)
			p = p.next
			continue
		}

		// 从[0,idx]随机选取一个值
		randIdx := rand.Intn(idx + 1)
		if randIdx < idx {
			result[randIdx] = p.val
		}
		p = p.next
	}

	return result
}

为了模仿n的无限大（即不知道n有多大），这里用链表实现，输入的是链表头节点

定义一个根据数组生成链表的函数

func New(data []int) *Node {
	if len(data) == 0 {
		return nil
	}

	root := &Node{
		val: data[0],
	}

	next := root

	for i := 1; i < len(data); i++ {
		next.next = &Node{
			val: data[i],
		}
		next = next.next
	}

	return root
}

测试

func main() {
	data := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
	for i := 1; i <= 11; i++ {
		head := New(data)
		result := Rand(head, i)
		fmt.Printf("n:%d,result:%v\n", i, result)
	}
}

测试代码定义了n=10的数组，i表示需要从数据源data中抽样的数据量，从1到11。

测试结果如下：

# 第一轮
n:1,result:[8]
n:2,result:[1 10]
n:3,result:[9 4 10]
n:4,result:[6 10 9 8]
n:5,result:[1 6 7 4 10]
n:6,result:[10 7 3 4 5 9]
n:7,result:[9 2 3 4 5 8 7]
n:8,result:[1 2 10 4 5 6 7 9]
n:9,result:[1 2 3 4 5 6 7 10 9]
n:10,result:[1 2 3 4 5 6 7 8 9 10]
n:11,result:[1 2 3 4 5 6 7 8 9 10]

# 第二轮
n:1,result:[7]
n:2,result:[1 8]
n:3,result:[9 7 4]
n:4,result:[1 10 8 4]
n:5,result:[1 9 8 4 10]
n:6,result:[1 10 8 4 7 6]
n:7,result:[10 2 3 4 5 9 8]
n:8,result:[1 2 10 4 5 6 7 8]
n:9,result:[1 2 10 4 5 6 7 8 9]
n:10,result:[1 2 3 4 5 6 7 8 9 10]
n:11,result:[1 2 3 4 5 6 7 8 9 10]