从可重复组合列表中获取含指定元素的索引及对应组合

背景说明 #

给定n=4（对象总数）、r=3（样本数）的非降序有序可重复组合（多重组合）列表，完整内容如下：

[1, 1, 1]
[1, 1, 2]
[1, 1, 3]
[1, 1, 4]
[1, 2, 2]
[1, 2, 3]
[1, 2, 4]
[1, 3, 3]
[1, 3, 4]
[1, 4, 4]
[2, 2, 2]
[2, 2, 3]
[2, 2, 4]
[2, 3, 3]
[2, 3, 4]
[2, 4, 4]
[3, 3, 3]
[3, 3, 4]
[3, 4, 4]
[4, 4, 4]

需求 #

无需生成完整组合列表（避免大内存占用），通过公式化方法直接获取：

1. 所有包含指定元素的组合的1-based索引
2. 对应的组合列表（保持与原列表一致的顺序，组合内元素非降序）

例如指定元素2时，预期返回：

• 1-based索引：2, 5, 6, 7, 11, 12, 13, 14, 15, 16
• 对应组合列表：

[[1, 1, 2], 
[1, 2, 2],
[1, 2, 3],
[1, 2, 4],
[2, 2, 2],
[2, 2, 3],
[2, 2, 4],
[2, 3, 3],
[2, 3, 4],
[2, 4, 4]]

核心实现方案 #

一、生成包含指定元素x的组合列表 #

多重组合的本质是「从n个元素中可重复选取r个的非降序排列」，对应数学模型为星与条问题：总组合数公式为C(n+r-1, r)（本例中C(4+3-1,3)=20，与列表长度一致）。

转化问题简化生成

包含至少1个元素x的组合，可转化为求解以下非负整数方程：

y₁ + y₂ + … + yₓ₋₁ + zₓ + yₓ₊₁ + … + yₙ = r-1

其中：

• yᵢ（i≠x）：元素i在组合中的出现次数
• zₓ = yₓ - 1：yₓ是元素x在组合中的出现次数（yₓ ≥1，因此zₓ ≥0）

组合构造步骤

1. 枚举上述方程的所有非负整数解；
2. 对每个解，构造组合：元素i（i≠x）出现yᵢ次，元素x出现zₓ+1次；
3. 将元素按非降序排列，得到最终组合；
4. 按「元素从小到大优先」的顺序枚举解，即可保证生成的组合顺序与原列表一致。

二、计算组合的1-based索引 #

对于非降序组合a₁≤a₂≤…≤aᵣ，需先将其映射为严格递增序列，再计算索引：

步骤1：映射为严格递增序列

对组合的每个元素做变换：bₖ = aₖ + k - 1（k=1~r），得到严格递增序列b₁<b₂<…<bᵣ（例如组合[1,2,2]对应b=[1,3,4]）。

步骤2：计算索引

索引 = 「严格小于该b序列的不重复组合数量」 + 1

计算小于b序列的组合数量的递推方法：

1. 初始化count=0，prev=0，remaining=r（剩余需要选取的元素个数）；
2. 遍历每个bₖ：
   • 累加i从prev+1到bₖ-1的组合数C(n+r-1 -i, remaining-1)到count；
   • 更新prev=bₖ，remaining -=1；
3. 最终索引为count+1。

示例验证（x=2, n=4, r=3） #

1. 组合生成：枚举方程y₁ + z₂ + y₃ + y₄ = 2的所有解，构造出的组合与示例完全一致；
2. 索引计算：
   • 组合[1,1,2]对应b=[1,2,4]，计算得小于它的组合数为1，索引=1+1=2；
   • 组合[2,2,2]对应b=[2,3,4]，计算得小于它的组合数为10，索引=10+1=11；
     所有索引结果与示例完全匹配。

这种方案无需生成完整组合列表，内存占用仅与目标组合数量相关，可高效处理大型n/r场景。

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