如何在Python/NumPy中高效生成固定维度的所有凸组合（和为1.0）

需要生成所有和为1.0的三维数组（凸组合），每个元素只能取[0.0, 0.2, 0.4, 0.6, 0.8, 1.0]中的值，例如[0.0,0.4,0.6]或[0.2,0.6,0.2]符合要求，[1.0,0.4,0.2]因和为1.6不符合。

当前通过生成全量组合再筛选的方式实现：

import numpy as np

ratios = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
result = np.stack(np.meshgrid(ratios, ratios, ratios), -1).reshape(-1, 3)
# 筛选和为1.0的组合
filtered = result[np.isclose(np.sum(result, axis=1), 1.0)]

但高维度场景下计算量极大——数十亿组合中仅不足千分之一满足条件，效率极低。

另有类似可变维度的问题（比如[1.0]或[0.2,0.2,0.2,0.2,0.2]均为有效解），但当前聚焦固定维度、固定和的场景，以下是更高效的实现方式：

核心优化思路 #

避免生成所有无效组合，通过整数转换简化问题，直接构造满足条件的数组：

1. 问题转换（规避浮点数精度） #

所有取值都是0.2的倍数，将每个元素乘以5，问题转化为：寻找三维非负整数数组，元素取值为[0,1,2,3,4,5]，且和为5（1.0×5=5）。既消除浮点数精度问题，又大幅降低计算量。

2. 固定三维的高效实现 #

直接遍历前两个维度的可能值，计算第三个维度是否符合要求，无需生成全量组合：

import numpy as np

def generate_3d_convex_combinations(target_sum=1.0, step=0.2):
    int_target = int(target_sum / step)
    int_values = {0, 1, 2, 3, 4, 5}  # 对应原浮点数乘以5的结果
    combinations = []
    
    for a in int_values:
        if a > int_target:
            continue
        for b in int_values:
            remaining = int_target - a - b
            if remaining in int_values and remaining >= 0:
                combinations.append([a, b, remaining])
    
    # 转换回原浮点数并去重
    return np.unique(np.array(combinations) * step, axis=0)

3. 通用固定维度实现（支持任意dim） #

如果需要支持更高固定维度，用回溯法生成符合条件的整数解，再转换回浮点数：

import numpy as np

def generate_fixed_dim_convex(dim, target_sum=1.0, step=0.2):
    int_target = int(target_sum / step)
    max_int_val = int(1.0 / step)  # 对应原取值的最大值1.0
    solutions = []
    
    def backtrack(current_dim, path, remaining):
        if current_dim == dim - 1:
            # 最后一个维度直接取剩余值，需在合法范围内
            if 0 <= remaining <= max_int_val:
                solutions.append(path + [remaining])
            return
        # 当前维度遍历可能的取值，从0到剩余值与最大值的较小值
        for val in range(0, min(max_int_val, remaining) + 1):
            backtrack(current_dim + 1, path + [val], remaining - val)
    
    backtrack(0, [], int_target)
    # 去重并转换回浮点数
    return np.unique(np.array(solutions) * step, axis=0)

效率对比 #

• 三维场景：原方法生成216个组合再筛选，新方法直接生成21个有效组合，计算量仅为原方法的10%左右
• 五维场景：原方法生成7776个组合，新方法仅生成约120个有效组合，效率提升数十倍

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