高效查找指定范围内数值多重集合的技术优化问询

核心思路 #

利用数组升序、单数组内元素差超过width的关键条件，结合多路归并+贪心分组+有序集合索引实现高效划分：

1. 用多路归并替代全局排序，以O(K log N)时间生成全局有序的元素流，同时保留每个元素的原坐标（数组索引+元素位置）。
2. 维护每个数组对应的可用集合有序列表：列表按集合的最小值升序排列，仅包含未使用该数组的集合，确保能快速找到符合宽度要求的候选集合。
3. 贪心匹配：对每个元素，优先加入第一个满足宽度要求且未使用其所属数组的集合；无匹配时新建集合，并将新集合加入其他数组的可用列表。

关键优化点 #

• 集合的最小值固定不变（因元素按升序加入，第一个元素即为最小值），因此可用列表可始终维持有序，二分查找有效，避免了之前均值更新导致的无序问题。
• 总集合数约为K/N（≈1e4），远小于元素总数，大幅降低了候选集合的遍历成本。
• 用布尔数组标记集合已使用的数组，O(1)时间判断是否可加入当前元素。

实现代码（核心逻辑） #

import heapq
from sortedcontainers import SortedList
import array

def partition_arrays(arrays, width):
    N = len(arrays)
    all_indices = list(range(N))
    
    # 多路归并初始化：堆中存储(元素值, 数组索引, 元素在数组中的位置)
    heap = []
    for i in range(N):
        if arrays[i]:
            heapq.heappush(heap, (arrays[i][0], i, 0))
    
    # 初始化集合列表与可用集合索引
    sets = []
    # available_sets[k]：存储未使用数组k的集合，格式为(集合最小值, 集合索引)，按最小值升序
    available_sets = [SortedList() for _ in range(N)]
    
    while heap:
        x_val, i, j = heapq.heappop(heap)
        # 将当前数组的下一个元素加入堆（如果存在）
        if j + 1 < len(arrays[i]):
            heapq.heappush(heap, (arrays[i][j+1], i, j+1))
        
        threshold = x_val - width
        found = False
        # 在可用集合中查找第一个满足宽度要求的集合
        idx = available_sets[i].bisect_left((threshold, -1))
        while idx < len(available_sets[i]):
            min_val, s_idx = available_sets[i][idx]
            if x_val - min_val <= width:
                # 加入该集合
                target_set = sets[s_idx]
                target_set['elements'].append((i, j))
                target_set['used'][i] = True
                # 从当前数组的可用列表中移除该集合（后续元素无法再加入）
                available_sets[i].pop(idx)
                found = True
                break
            idx += 1
        
        if not found:
            # 新建集合
            new_set = {
                'min_val': x_val,
                'used': array.array('b', [False]*N),  # 用array节省内存
                'elements': [(i, j)]
            }
            new_set['used'][i] = True
            s_idx = len(sets)
            sets.append(new_set)
            # 将新集合加入其他数组的可用列表
            for k in all_indices:
                if k != i:
                    available_sets[k].add((x_val, s_idx))
    
    # 返回每个集合的元素原坐标
    return [s['elements'] for s in sets]

无第三方库适配方案 #

若无法使用sortedcontainers，可改用bisect维护普通列表，结合标记法处理无效集合：

import heapq
import bisect
import array

def partition_arrays_no_lib(arrays, width):
    N = len(arrays)
    all_indices = list(range(N))
    
    heap = []
    for i in range(N):
        if arrays[i]:
            heapq.heappush(heap, (arrays[i][0], i, 0))
    
    sets = []
    # available_sets[k]：存储(集合最小值, 集合索引, 是否有效)，按最小值升序
    available_sets = [[] for _ in range(N)]
    
    while heap:
        x_val, i, j = heapq.heappop(heap)
        if j + 1 < len(arrays[i]):
            heapq.heappush(heap, (arrays[i][j+1], i, j+1))
        
        threshold = x_val - width
        found = False
        lst = available_sets[i]
        # 二分查找第一个符合阈值的位置
        idx = bisect.bisect_left(lst, (threshold, -1, True))
        
        while idx < len(lst):
            min_val, s_idx, valid = lst[idx]
            if not valid:
                idx += 1
                continue
            if x_val - min_val <= width:
                target_set = sets[s_idx]
                target_set['elements'].append((i, j))
                target_set['used'][i] = True
                # 标记该集合对当前数组无效
                lst[idx] = (min_val, s_idx, False)
                found = True
                break
            idx += 1
        
        if not found:
            new_set = {
                'min_val': x_val,
                'used': array.array('b', [False]*N),
                'elements': [(i, j)]
            }
            new_set['used'][i] = True
            s_idx = len(sets)
            sets.append(new_set)
            entry = (x_val, s_idx, True)
            # 加入其他数组的可用列表
            for k in all_indices:
                if k != i:
                    bisect.insort(available_sets[k], entry)
    
    return [s['elements'] for s in sets]

性能说明 #

• 时间复杂度：O(K log N + K log M + N*M log M)，其中M为总集合数（≈1e4），适配千万级元素量。
• 内存占用：集合标记数组单集合仅需1KB，总集合内存约10MB；元素坐标存储约80MB，整体内存可控。

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