等概率采样最大值的索引
今天，我参加了一场面试，一面竟然全是考算法。真是崩溃。
如题，就是我遇到的第一题，直接给我难住了。题目意思如下：
给定一个整数数组，设计一个方法，返回数组中最大值的下标。如果，最大值有重复出现，则按照相等的概率返回任意一个最大值的下标。附加条件：设计一个时间复杂度 O(n)，空间复杂度 O(1) 的算法。
这道题我在 LC 没找到，在 SO 倒是搜到了。—> 传送门
这个题目非常的抽象，我们来举两个例子，比如，给定一个数组是 [1,2,3,3,3]，这个数组的最大值是 3，其下标是 2，3 和 4，那么，你的方法要返回 2，3，4 中的一个，出现概率是三分之一。
再举一个例子，输入数组是[1,2,3,4,5]，最大值是 5，其下标是 4，因为只出现了一次，那么你的方法，每次返回值都是 4，也即 100% 概率返回 4。
反复思考了几遍，我才理解了这题目的意思。我的第一个思路很简单，就是用一个数组，记录下最大值的所有下标，然后用一个随机数取整，然后挑一个下标返回。
代码大概可以这样写：1
2
3
4
5
6
7
8
9
10
11
12
13
14def findMaxIndex(arr: List[int])->int:
maxVal = arr[0]
indices = [0]
if len(arr) == 1: return 0
for i in range(1, len(arr)):
if arr[i] == maxVal:
indices.append(i)
elif arr[i] > maxVal:
maxVal = arr[i]
indices = [i]
# choice 是一个库函数，作用是随机挑一个元素返回
return random.choice(indices)
当然，面试的时候我只是口述了这个思路，根本就没写，因为这个算法明显不满足要求。最大值出现的频次是不确定的，也可能直接就出现 n 次。所以，存储最大值的下标的数组，明显破坏了空间复杂度约束。（上面的代码是我回家后，反思的时候写的，大家看最后一行，用了一个叫 choice 的方法，从一个数组里，随机返回一个元素，其实正是等概率返回其中一个元素的意思。当我们可以生成一个 [0,1] 的随机数的时候，怎样才能实现一个 choice 方法呢？毛想想好像很容易，真想写的时候，才发现我并不会写这个方法，而且，很奇妙的时候，如果我会写这个 choice 方法，那么做出这个题目也就没什么难的了，大家可以想想，如果题目换成实现一个 Python random 包的 choice 方法，你会怎么实现？）
因为我是第一次看到这个题目，又是在面试的场景下，我真的大脑一片空白，虽然我已经很努力压榨我的大脑了，仍然是一片空白。只好厚颜无耻地要求面试官给出提示。面试不是竞赛，可以厚颜无耻。
如果读者朋友你也没有见过这个题目，你可以现在假设一下，面试官怎样提示你，你才能想到正确的写法？反正我现在仍然是懵，我如果是面试官，我都不知道怎么去提示候选人，如果她/他真的没概念的话。
我的面试官是这么说的，你先考虑一个最简单的情况，[1,2,3]，这个数组，你会返回 2（唯一最大值 3 的下标），这时候，又多了一个数字 3，变成了 [1,2,3,3]，这个时候你会怎么办？
妈呀！这也叫提示啊？我听完了，完全是十脸懵逼好不好。不过，懵逼也不能慌，要装白小纯，单就这个情况来看，不考虑以后，现在只有两个 3，我应该按照 50% 的概率去挑选 2 和 3 两个下标，没错吧。所以，我可以生成个随机数，如果小于 0.5 我就返回 2，大于 0.5 我就返回 3。很朴实嘛。
后面就不知道了。继续厚颜无耻。面试官说，还是这个例子，这时候，又来了一个 3，变成 [1,2,3,3,3]，这时候你怎么办？我还能怎么办，我还绕在第一个思路里出不来呢，因为我刚才已经返回了 2，或者 3，现在又来了一个怎么办呢？我想，他这么提示，莫不是让我用递归思想？不不不，我肯定递归不出来。假设我刚才选定了 2，现在怎么办？假设我刚才选了 3，现在又怎么办？我只选了一个，另一个被我放弃了啊，又多了一个怎么办呢？
等等，等等，我刚才已经按照 50% 的概率选中了 2 或者 3，也就是，我上一次选定的数字，既可能是 2，也可能是 3，虽然我选定了一个，但是给另一个的机会也是均等的啊，就完全没必要记住另一个是几了啊！！！到这里我已经开窍了，终于被提示出来了！太不容易了。
唯一的难点来了，现在又多了一个 3，其下标是 4，直觉上，我觉得，只要把刚才 50% 的概率给降下来，降到 33%，然后把 4 这个下标混进去，就可以了。所以，应该用 66.6% 的概率选择上一次的选择，按照 33.3% 的概率选择 4，是不是就可以了？
其实，现在我已经知道这就是正确答案了，终于两次厚颜无耻地要求提示下，我想出来了。当然，现场我是非常不确定的，也没什么自信，毕竟太烧脑了，我状态也不轻松。好在面试官没有为难我，说这是对的，用数学归纳法不难证明。你把这思路写出来写成代码吧……1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18def findMaxIndex(arr: List[int]) -> int:
maxVal = arr[0]
selected = 0
count = 1
if len(arr) == 1: return 0
for i in range(1, len(arr)):
if arr[i] == maxVal:
# 下面这行很重要，注意！
if random.random() <= 1 / (count + 1):
selected = i
count += 1
elif arr[i] > maxVal:
selected = i
count = 1
return selected
其实我还是懵的。就硬着头皮开始写，写得多难看我就不提了。不过好歹写出来了。其中生成随机数，并且控制概率输出的部分，我大脑一片空白，胡乱写了个语法完全不对的，号称伪代码的东西，就蒙混了，面试官说，能说清楚意思就行了。
注意看第 11 行，假设我们已经遇到了 n 个数字，并且我们按照 1/n 的概率，选中了一个下标 a 的时候，这时候，又来了一个数字，第 n + 1 个，其下标是 b，那么，这时候我们到底是返回 a 还是返回 b 呢？注意，我们最终返回的下标，其出现的概率是 1/(n+1)，我们已经知道 a 出现的概率是 1/n，这时候怎么处理 b 呢？
我们应该按照 1/(n+1) 的概率返回 b，并且按照 n/(n+1) 的概率返回 a。这样，返回 a 的概率是 1/n × n/(n+1) 正好等于 1/(n+1)，这样就可以满足题目的要求了。简直完美。这个算法，也叫蓄水池采样（Reservoir sampling）算法。非常巧妙，用一个迭代的算法，完成了等概率的分配。
这时候，35 分钟过去了，面试官说，还剩 10 分钟，要不再来一道吧……我靠！！我内心是崩溃的，因为刚逃脱一个危险，又来，实在是压力倍增。幸亏第二题遇到了我做过的题目，著名的打家劫舍。可是我两个多月没刷题了啊，就算一直在刷也不能总是在刷动态规划吧。还好他说时间不够，就讲解思路，简直如蒙大赦。这个打家劫舍也不是最平凡的那个，而是环形打家劫舍，虽然我曾经做过，但是当时真的吃得不透，这时候，只好信口胡诌，好在原理是没差别的，就在于怎么处理环形。被我蒙混过去了。我暗忖，真要写，怕是没那么容易写对的。
虽然我不知道结果如何，但是面试结束后，我还是有点小成就感的，竟然还是答出来了，今年刷了三个多月的题目，真的没白刷。可见还是熟能生巧的一个过程。虽然我又闲散了两个月，已经又生疏了，但是和以前真的不同了。这次面试无论成败我想我都可以泰然面对了。并且涌起了新的刷题的动力。还是要持之以恒啊。