如何高效确定长度为k的整数n分拆的秩（索引）？

问题背景 #

需要计算整数n的长度为k的分拆在逆字典序中的索引，无需生成所有分拆。例如n=10、k=3时，逆字典序的分拆列表如下：

0 (8, 1, 1)
1 (7, 2, 1)
2 (6, 3, 1)
3 (6, 2, 2)
4 (5, 4, 1)
5 (5, 3, 2)
6 (4, 4, 2)
7 (4, 3, 3)

目标是快速获取特定分拆（如[5,3,2]）对应的索引。

核心思路 #

逆字典序的分拆按非递增排列，且前面的元素越大越靠前。计算目标分拆的索引，本质是统计所有排在它前面的分拆总数，可以通过递归累加符合条件的分拆数量实现，核心依赖你编写的num_partitions(n, k, p)函数（修正后）。

修正后的辅助函数 #

你的num_partitions(n, k, p)函数用于计算「总和为n、分为k部分、每部分不超过p的分拆数」，需修正函数调用的名称错误：

def num_partitions(n, k, p):
    # 边界条件：无效输入返回0
    if n < 0 or k < 0 or p <= 0:
        return 0
    # 恰好分完的情况返回1
    if n == 0 and k == 0:
        return 1
    # 两种情况的和：选一个p作为部分，剩下的n-p分k-1部分且每部分不超过p；不选p，即最大部分不超过p-1
    return num_partitions(n - p, k - 1, p) + num_partitions(n, k, p - 1)

索引计算函数 #

定义递归函数计算目标分拆的索引，分为全局计算和受限计算（限制后续元素的最大值）：

def get_partition_index(partition, n, k):
    # 确保分拆是非递增排列
    assert sorted(partition, reverse=True) == partition, "分拆必须按非递增排列"
    assert sum(partition) == n and len(partition) == k, "分拆不符合n和k的要求"
    
    def restricted_index(sub_partition, sub_n, sub_k, max_val):
        if sub_k == 1:
            return 0
        current = sub_partition[0]
        count = 0
        # 当前位置元素的最大值：不能超过max_val，且剩下的sub_k-1部分至少为1，所以最大是sub_n - (sub_k -1)
        upper = min(max_val, sub_n - (sub_k - 1))
        # 累加所有当前元素大于current的分拆数
        for m in range(current + 1, upper + 1):
            count += num_partitions(sub_n - m, sub_k - 1, m)
        # 递归处理剩余部分
        count += restricted_index(sub_partition[1:], sub_n - current, sub_k - 1, current)
        return count
    
    return restricted_index(partition, n, k, n)

验证例子 #

以分拆[5,3,2]、n=10、k=3为例：

1. 首先统计第一个元素大于5的分拆数：
   • m=6：num_partitions(4,2,6)=2（对应分拆(6,3,1)、(6,2,2)）
   • m=7：num_partitions(3,2,7)=1（对应分拆(7,2,1)）
   • m=8：num_partitions(2,2,8)=1（对应分拆(8,1,1)）
     这部分总和为2+1+1=4。
2. 然后处理剩余分拆[3,2]，剩余总和5、剩余部分数2、最大元素限制为5：
   • 统计当前元素大于3的分拆数：m=4时num_partitions(1,1,4)=1（对应分拆(5,4,1)的剩余部分(4,1)）
   • 递归处理最后一个元素，返回0
     这部分总和为1+0=1。
3. 总索引为4+1=5，与例子完全一致。

优化建议 #

如果需要处理较大的n和k，可以给num_partitions添加记忆化（缓存），避免重复计算，提升效率：

from functools import lru_cache

@lru_cache(maxsize=None)
def num_partitions(n, k, p):
    if n < 0 or k < 0 or p <= 0:
        return 0
    if n == 0 and k == 0:
        return 1
    return num_partitions(n - p, k - 1, p) + num_partitions(n, k, p - 1)

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