如何高效实现多类型数值到非重叠Bucket的快速分配？

问题背景 #

我编写了一段将数值分配至对应Bucket的代码：现有N个分属Y种类型的Bucket，每个类型对应一个数值，需判断数值所属Bucket（数值≥Bucket起始值且≤Bucket终止值时归属该Bucket）。当前要查询用户ID5607、线程ID11、服务器ID3的归属Bucket。原实现的时间复杂度为O(N)，但代码会被频繁调用，需要更高效的方案。

原代码与输出 #

原代码 #

from operator import itemgetter
buckets2 = {
  "buckets": [
    {"type": "user", "name": "users-bucket1", "start": 0, "end": 5000},
    {"type": "user", "name": "users-bucket2", "start": 5001, "end": 10000},
    {"type": "user", "name": "users-bucket3", "start": 10001, "end": 15000},
    {"type": "server", "name": "server-bucket1", "start": 0, "end": 3},
    {"type": "server", "name": "server-bucket2", "start": 4, "end": 7},
    {"type": "server","name": "server-bucket3", "start": 8, "end": 10},
    {"type": "thread", "name": "threads-bucket1", "start": 0, "end": 4},
    {"type": "thread", "name": "threads-bucket2", "start": 5, "end": 9},
    {"type": "thread", "name": "threads-bucket3", "start": 10, "end": 14}
    ]
 }

buckets2["buckets"].sort(key=itemgetter("start"), reverse=True)

def route2(routing_information):
    buckets = {}
    for item in buckets2["buckets"]:
       if item["end"] >= routing_information[item["type"]]:
           buckets[item["type"]] = item["name"]
    return buckets

print(route2({ "server": 3, "user": 5607, "thread": 11}))

原输出 #

{'user': 'users-bucket2', 'thread': 'threads-bucket3', 'server': 'server-bucket1'}

优化方案 #

核心思路 #

将Bucket按类型分组，对每个类型的Bucket列表排序后，使用二分查找替代全量遍历。预处理仅需执行一次，单次查询的时间复杂度从O(N)降至O(Y log M)（Y为类型数量，M为单类型Bucket的数量），大幅提升高频调用场景下的性能。

优化代码实现 #

from bisect import bisect_right

# 预处理：按类型分组并排序，仅需执行一次
bucket_groups = {}
for bucket in buckets2["buckets"]:
    typ = bucket["type"]
    if typ not in bucket_groups:
        bucket_groups[typ] = []
    bucket_groups[typ].append(bucket)

# 对每个类型的Bucket按start升序排序，并提取end值列表用于二分查找
for typ in bucket_groups:
    # 按起始值排序，确保区间连续不重叠（假设原数据符合该规则）
    bucket_groups[typ].sort(key=lambda x: x["start"])
    # 单独存储end值列表，避免二分查找时重复访问字典
    bucket_groups[typ].append({
        "end_list": [b["end"] for b in bucket_groups[typ]]
    })

def optimized_route(routing_info):
    result = {}
    for typ, value in routing_info.items():
        bucket_list = bucket_groups[typ]
        end_list = bucket_list[-1]["end_list"]
        # 用bisect_right找到第一个大于value的end值索引，前一个即为目标Bucket
        idx = bisect_right(end_list, value) - 1
        result[typ] = bucket_list[idx]["name"]
    return result

# 测试优化后的函数
print(optimized_route({"server": 3, "user": 5607, "thread": 11}))

关键说明 #

1. 预处理逻辑：按类型分组后排序，确保每个类型的Bucket区间是有序的，为二分查找提供基础。
2. 二分查找效率：bisect_right是Python内置的高效二分实现，时间复杂度为O(log M)，远快于全量遍历。
3. 适用前提：假设每个类型的Bucket区间是连续、不重叠且能覆盖所有查询数值（原数据已满足该条件）。

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