如何将筛选符合差值条件的函数优化至O(nlogn)时间复杂度？

需求说明：给定一个数字列表，找出能选取的最大数字数量，要求选取的任意两个数字的差值不超过t。当前实现的时间复杂度在最坏情况下为O(n²)，需要优化至O(nlogn)。

当前代码分析 #

当前代码首先对列表排序（O(nlogn)），然后通过嵌套循环遍历每个起始位置，统计满足条件的元素数量。但内层循环在最坏场景（如所有元素差值都<=t）下会遍历整个剩余数组，导致总时间复杂度变为O(n²)，无法达到目标。

def find_numbers(num_list, t):
    sorted_list=sorted(num_list)
    counter=0
    k=0
    n=0
    for i in range(len(sorted_list)):
        for j in range(k, len(sorted_list)):
            if sorted_list[j]-sorted_list[k]<=t:
                counter+=1
            else:
                break
        k+=1
    
        if counter>n:
            n=counter
        counter=0
    return n 

运行示例：

print(find_numbers([2, 7, 14, 11, 7, 15], 11)) # 5
print(find_numbers([4, 2, 7, 1], 0)) # 1
print(find_numbers([7, 3, 1, 5, 2], 2)) # 3

第三个示例说明：从列表[7,3,1,5,2]中可选取3个数字：3、1和2，这些数字之间的差值均不超过2。

优化方案1：滑动窗口（双指针）法 #

利用数组已排序的特性，使用两个指针维护一个窗口，窗口内的元素满足最大值-最小值 <=t。由于数组递增，左指针右移时，右指针无需重置，只需继续向右扩展即可，整个遍历过程为O(n)，加上排序的O(nlogn)，总时间复杂度为O(nlogn)。

优化后代码 #

def find_numbers(num_list, t):
    sorted_list = sorted(num_list)
    max_count = 0
    right = 0
    n = len(sorted_list)
    for left in range(n):
        # 扩展右指针，直到窗口内差值超过t
        while right < n and sorted_list[right] - sorted_list[left] <= t:
            right += 1
        # 当前窗口的元素数量是right - left
        current_count = right - left
        if current_count > max_count:
            max_count = current_count
    return max_count

逻辑说明 #

1. 先对列表排序，这是O(nlogn)的核心操作。
2. 初始化右指针right为0，max_count记录最大符合条件的元素数量。
3. 遍历每个左指针left：
   • 不断右移right，直到sorted_list[right] - sorted_list[left] > t，此时窗口[left, right-1]内的所有元素都满足差值<=t。
   • 计算当前窗口的元素数量right - left，更新max_count。
4. 最终返回max_count。

优化方案2：二分查找法 #

对于每个左边界left，利用二分查找找到第一个大于sorted_list[left] + t的元素位置，该位置与left的差值就是当前左边界下的最大元素数量。每个二分查找是O(logn)，n次查找加上排序的O(nlogn)，总时间复杂度也是O(nlogn)。

代码示例 #

import bisect

def find_numbers(num_list, t):
    sorted_list = sorted(num_list)
    max_count = 0
    n = len(sorted_list)
    for left in range(n):
        # 找到第一个大于sorted_list[left]+t的索引
        right_idx = bisect.bisect_right(sorted_list, sorted_list[left] + t)
        current_count = right_idx - left
        if current_count > max_count:
            max_count = current_count
    return max_count

验证结果 #

两种优化后的代码都能通过给定的测试用例：

• 输入[2, 7, 14, 11, 7, 15]和t=11，返回5；
• 输入[4, 2, 7, 1]和t=0，返回1；
• 输入[7, 3, 1, 5, 2]和t=2，返回3。

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