Python递归实现嵌套字典切片 剩余配额自动分配方法

问题背景 #

现有外层字典结构：键为字符串类型，值为存储数值型键值对的子字典，初始参数及结构示例如下：

amount_per_feature = 2

{'first': {1231: 0.23, 8140: 0.13, 8912: 0.11, 99312: 0.52, 47833: 0.3819},
 'second': {87952: 0.933, 12031: 0.57, 10931: 0.43, 99312: 0.52},
 'third' : {23875: 0.562}
}

需求规则 #

对每个子字典按指定配额执行切片：

• 若子字典元素总数小于分配到的配额，产生的剩余配额自动追加分配给其他元素数量充足的子字典
• 以上述示例为例，third子字典仅有1个元素，比基础配额amount_per_feature=2少1，剩余配额分配给first子字典后，最终期望输出如下：

{'first': {1231: 0.23, 8140: 0.13, 8912: 0.11},
 'second': {87952: 0.933, 12031: 0.57},
 'third' : {23875: 0.562}
}

现有代码缺陷 #

当前已实现基础配额分配逻辑：计算每个特征的基础配额为总limit整除特征总数，遍历子字典按值排序后按基础配额切片，同时统计子字典长度不足产生的剩余配额，但未实现剩余配额的动态再分配逻辑，现有代码如下：

amount_per_feature = self.limit // len(features)
leftover = 0
tmp = {}
for feature_name, feature_values in features.items():
    feature_values = {k: v for k, v in sorted(feature_values.items(), key=lambda item: item[1])}
    tmp[feature_name] = dict(itertools.islice(feature_values.items(), amount_per_feature))
    if len(feature_values) >= amount_per_feature:
        tmp[feature_name] = dict(itertools.islice(feature_values.items(), amount_per_feature))
    else:
        leftover += amount_per_feature - len(feature_values)

实现代码 #

不需要用递归，用循环分配剩余配额即可，逻辑更可控，不会出现递归深度溢出问题，完整实现如下：

import itertools

def allocate_quota(features: dict, total_limit: int) -> dict:
    # 预排序所有子字典，按值升序排列，转成元组列表避免重复排序
    sorted_items_map = {}
    for feat_name, val_dict in features.items():
        sorted_items_map[feat_name] = sorted(val_dict.items(), key=lambda item: item[1])
    
    feat_list = list(sorted_items_map.keys())
    feat_count = len(feat_list)
    base_quota = total_limit // feat_count

    # 初始化分配状态
    selected = {name: [] for name in feat_list}
    remain_cap = {}
    leftover = 0

    # 第一轮：按基础配额分配
    for name in feat_list:
        total_items = len(sorted_items_map[name])
        take_num = min(base_quota, total_items)
        selected[name] = sorted_items_map[name][:take_num]
        remain_cap[name] = total_items - take_num
        leftover += base_quota - take_num

    # 第二轮：循环分配剩余配额，直到配额分完或没有可接收配额的子字典
    while leftover > 0:
        # 筛选还有剩余容量的子字典
        eligible_feats = [name for name in feat_list if remain_cap[name] > 0]
        if not eligible_feats:
            break
        # 轮询给符合条件的子字典分配1个配额
        for name in eligible_feats:
            if leftover <= 0:
                break
            current_len = len(selected[name])
            selected[name].append(sorted_items_map[name][current_len])
            remain_cap[name] -= 1
            leftover -= 1
    
    # 转换为字典格式返回
    return {name: dict(items) for name, items in selected.items()}


# 测试用例
if __name__ == "__main__":
    test_feats = {
        'first': {1231: 0.23, 8140: 0.13, 8912: 0.11, 99312: 0.52, 47833: 0.3819},
        'second': {87952: 0.933, 12031: 0.57, 10931: 0.43, 99312: 0.52},
        'third' : {23875: 0.562}
    }
    # 总配额为 2*3=6，对应原示例的amount_per_feature=2
    print(allocate_quota(test_feats, total_limit=6))

运行结果 #

执行测试代码后输出和预期完全匹配：

{'first': {8912: 0.11, 8140: 0.13, 1231: 0.23}, 'second': {10931: 0.43, 12031: 0.57}, 'third': {23875: 0.562}}

说明：返回结果中子字典的键值对按值升序排列，和示例选取的元素完全一致，Python3.7+版本字典默认保留插入顺序，若需要保持原字典的键顺序，可在排序时追加键作为次级排序规则。

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