在面美团实习的时候,面试官问了这道蓄水池问题:给定一个链表,长度未知(但是大于k),你只能遍历一次,要求随机挑出k个节点。
刚听完题目的我是一脸蒙蔽的,what is the fuck? 长度未知?遍历一遍?那每个节点被挑选的概率如何计算?
我当时的答案是遍历两遍链表。
后来想想这样的确很蠢,蓄水池问题的一个变种:有一个网络数据流,要你随机抓一个包,所有包被抓到的概率要求相等。
如果用我的遍历两遍链表的方法肯定不行!
然后我google了一下,发现了神奇的蓄水池解法:
对于链表随机挑k个节点,先把前k个节点作为已经选定的节点,从第k+1个节点开始,每个节点以1/(k+1)的概率被选中,并从选定的节点中随机挑选一个替换它,成为新的候选节点。
vector<int> vec(k,0); int num = k,curlen = k+1; while(num--){ vec[k-num] = head->val; head = head->next; } while(head){ int m = rand()%curlen; if(m < k) swap(vec[m],head->val); head = head->next; curlen++; }
这样一来,第i个节点被选中的概率为1/i。而它之前的节点i-1被选中的概率为 1/(i-1)*(1-1/i) = 1/i。满足题意。这个思路实在是666。