如何高效获取数字列表中元素两两间距不小于d的k元子集

问题原因 #

你原方案使用itertools.combinations枚举所有组合的时间复杂度为组合级O(C(n,k))，当n=50、k=20时组合数高达4.7e13，完全无法遍历完成，所以会出现运行过慢的问题。

高效解法 #

使用贪心算法即可解决，时间复杂度仅为O(nlogn)，可以轻松处理千级以上规模的输入，同时可以直接判断合法子集是否存在。

算法逻辑

1. 先将输入的数字列表按升序排序
2. 从左到右遍历排序后的列表，每次尽可能早地选中符合条件的元素：第一个元素直接选中，后续每个元素只要和上一个选中元素的差值≥d，就将其加入结果集
3. 遍历完成后，如果结果集的长度≥k，说明存在符合要求的子集，取结果集的前k个元素返回即可；如果长度不足k，说明不存在符合要求的子集，返回None

代码实现

def k_subset_at_least_d_apart(nums, k, d):
    sorted_nums = sorted(nums)
    res = []
    last = float('-inf')
    for num in sorted_nums:
        if num - last >= d:
            res.append(num)
            last = num
            # 提前终止：选够k个直接返回，无需继续遍历
            if len(res) == k:
                return tuple(res)
    return None

测试验证

lst = [1, 3, 6, 10, 15, 20, 27, 28, 30, 40]
subset = k_subset_at_least_d_apart(lst, k=5, d=6)
print(subset)
# 输出：(1, 10, 20, 27, 40)

正确性说明 #

贪心策略可以选出当前数组中最多的符合间隔要求的元素，如果最多可选数量都达不到k，就不存在合法子集；如果能达到，选中的前k个元素天然满足两两差值≥d的要求。

内容的提问来源于stack exchange，提问作者Shaun Han