如何用闭式函数将itertools组合映射为对应索引？

我在使用Python的itertools.combinations()生成组合时，希望能高效地通过生成的有序元组直接映射到它对应的枚举索引。比如用以下代码生成索引和组合：

import itertools

it = itertools.combinations(range(6), 4)
for i in enumerate(it):
    print(i)

输出结果：

(0, (0, 1, 2, 3))
(1, (0, 1, 2, 4))
(2, (0, 1, 2, 5))
(3, (0, 1, 3, 4))
(4, (0, 1, 3, 5))
(5, (0, 1, 4, 5))
(6, (0, 2, 3, 4))
(7, (0, 2, 3, 5))
(8, (0, 2, 4, 5))
(9, (0, 3, 4, 5))
(10, (1, 2, 3, 4))
(11, (1, 2, 3, 5))
(12, (1, 2, 4, 5))
(13, (1, 3, 4, 5))
(14, (2, 3, 4, 5))

我需要能快速把元组(0, 3, 4, 5)映射到索引9，并且这个方法要能通用，处理大规模组合时也要高效——之前尝试的多索引方案速度太慢。

我已经有个思路起点：首个数字d₁的出现次数可以表示为C(n-(d₁+1), p-1)（比如n=6、p=4时，0出现C(5,3)=10次，1出现C(4,3)=4次），但不知道怎么把这个思路推广到后续的d₂、d₃、d₄，也不清楚如何转换成最终的索引。

通用闭式函数实现思路 #

这个问题本质是计算字典序下，当前组合之前有多少个组合，总和就是当前组合的索引。核心利用组合数的性质，分步累加每个位置上“更小的合法组合数量”：

对于长度为p的组合元组(d₀, d₁, ..., d_{p-1})（满足0 ≤ d₀ < d₁ < ... < d_{p-1} < n），索引的计算步骤如下：

1. 初始化索引为0，前一个元素的初始值设为prev = -1（表示第一个元素之前的虚拟边界）。
2. 遍历组合中的每个元素d_i（从第0个到第p-1个）：
   • 计算当前位置上，比d_i小且大于prev的元素能构成的组合数：这部分等于从prev+1到d_i-1的每个数x，对应剩下p - i - 1个元素从x+1到n-1中选取的组合数之和，可简化为组合数的差值：comb(n - prev - 1, p - i) - comb(n - d_i, p - i)。
   • 将这个差值加到索引上。
   • 更新prev为当前的d_i。
3. 遍历结束后，得到的索引就是目标值。

代码实现 #

使用Python 3.10+自带的math.comb函数（如果版本更低，可以自己实现组合数计算函数）：

import math

def combination_to_index(comb_tuple, n):
    p = len(comb_tuple)
    index = 0
    prev = -1
    for i, d in enumerate(comb_tuple):
        # 计算当前位置的贡献
        index += math.comb(n - prev - 1, p - i) - math.comb(n - d, p - i)
        prev = d
    return index

# 测试示例
print(combination_to_index((0, 3, 4, 5), 6))  # 输出9
print(combination_to_index((2, 3, 4, 5), 6))  # 输出14

原理验证 #

以(0, 3, 4, 5)为例：

• 第0个元素d=0，prev=-1：comb(6 - (-1) -1,4-0) - comb(6-0,4-0) = comb(6,4)-comb(6,4)=15-15=0，索引加0，prev更新为0。
• 第1个元素d=3，prev=0：comb(6-0-1,4-1) - comb(6-3,4-1) = comb(5,3)-comb(3,3)=10-1=9，索引加9，现在索引是9，prev更新为3。
• 第2个元素d=4，prev=3：comb(6-3-1,4-2) - comb(6-4,4-2)=comb(2,2)-comb(2,2)=1-1=0，索引加0，prev更新为4。
• 第3个元素d=5，prev=4：comb(6-4-1,4-3) - comb(6-5,4-3)=comb(1,1)-comb(1,1)=1-1=0，索引加0。
  最终索引为9，和预期一致。

这种方法的时间复杂度是O(p)，p是组合的长度，处理大规模组合时效率远高于遍历生成所有组合的方案。

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