今天分享一道 Cresta 面试真题，题目考察的是在机器学习数据集中，如何根据每个类别的数据量，计算一个尽可能均衡的采样方案。
题目英文原文整理
You have a dataset which contains examples labeled with different classes, for example an animal image classification dataset containing different types of animals, one animal per image.
You want to randomly sample an n-element subset from this dataset in such a way that the classes are as balanced as possible, i.e. equal or with the least difference between the number of examples sampled from each class.
Write a function which, given n and the number of examples in each class, computes how many examples should be sampled from each class.
给定一个分类数据集，每个类别有不同数量的样本。
现在希望随机抽取总共 n 个样本，并且让各个类别被抽到的数量尽可能平均。
函数输入：
n = 需要抽取的总样本数
class_counts = 每个类别当前拥有的样本数量函数输出：
每个类别应该抽取多少个样本需要注意：某些类别本身样本数可能很少，不能超过该类别已有数量。
例如：
n = 8
class_counts = [2, 10, 10]最均衡的结果不是 [3, 3, 2]，因为第一个类别最多只有 2 个样本。
所以结果应该是：
[2, 3, 3]解题思路
这道题本质上是一个 带容量上限的均分问题。
如果所有类别样本都足够多，那么直接平均即可：
n // num_classes多出来的样本再分给一部分类别。
但问题在于，有些类别的样本数量可能不足。例如某个类别只有 1 个样本，就不能强行分配 3 个。
所以正确做法是：
- 先尽量让所有类别一起“涨水位”；
- 如果某个类别达到自己的最大容量，就固定它；
- 剩下的样本继续在其他还没满的类别中平均分配；
- 最终保证总数等于
n，并且各类别差距尽可能小。
这类似一个 water filling（水位填充） 算法。
复杂度分析
排序需要：
O(k log k)其中 k 是类别数量。
后续扫描每个类别一次，因此整体复杂度是：
O(k log k)空间复杂度是：
O(k)面试考点总结
这道题看起来像普通的随机采样题，但真正考察的是：
- 能否意识到每个类别有容量上限；
- 能否保证总采样数正好等于
n； - 能否让类别之间尽可能均衡；
- 能否处理小类别样本不足的情况；
- 能否写出比逐个样本分配更高效的算法。
实际随机抽样时，先用这个函数算出每个类别应该抽多少个，然后再分别从每个类别里随机抽取对应数量即可。
我们也有代面试，面试辅助，OA代写等服务助您早日上岸~
