如何无过滤、无类型修正生成满足约束的DataFrame唯一组合

需求说明 #

基于布尔型（仅含True/False）和浮点型原始数组，生成包含所有合法唯一组合的DataFrame，需满足：

• 内存限制：禁止先生成全量组合再过滤无效行
• 约束规则：每行仅允许一个浮点列有非零值，其余浮点列必须为0
• 现有方案痛点：需要事后过滤无效行，还要手动修正布尔列类型，效率低下

示例输入 #

inputs = {
    "a": [True, False],
    "b": [True, False],
    "c": [0.0, 0.1, 0.2],
    "d": [0.0, 0.1, 0.2, 0.3],
}

bool_inputs = {"a", "b"}
float_inputs = {"c", "d"}

示例输出 #

a      b    c    d
0    True   True  0.0  0.0
1    True  False  0.0  0.0
2   False   True  0.0  0.0
3   False  False  0.0  0.0
4    True   True  0.1  0.0
5    True  False  0.1  0.0
6   False   True  0.1  0.0
7   False  False  0.1  0.0
8    True   True  0.2  0.0
9    True  False  0.2  0.0
10  False   True  0.2  0.0
11  False  False  0.2  0.0
12   True   True  0.0  0.1
13   True  False  0.0  0.1
14  False   True  0.0  0.1
15  False  False  0.0  0.1
16   True   True  0.0  0.2
17   True  False  0.0  0.2
18  False   True  0.0  0.2
19  False  False  0.0  0.2
20   True   True  0.0  0.3
21   True  False  0.0  0.3
22  False   True  0.0  0.3
23  False  False  0.0  0.3

现有代码（存在内存浪费问题） #

import numpy as np
import pandas as pd


input_arrays = list(inputs.values())
results = np.array(np.meshgrid(*input_arrays)).T.reshape(-1, len(inputs))

df = pd.DataFrame(results, columns=list(inputs.keys()))

df[list(bool_inputs)] = df[list(bool_inputs)].astype(bool)
df = df[~(df[list(float_inputs)] > 0).all(axis=1)]
df = df.reset_index(drop=True)

优化实现方案 #

核心思路 #

1. 先生成布尔列的全量合法组合（布尔列无约束，所有组合都有效）
2. 分块生成浮点列的合法组合：
   • 第一块：所有浮点列均为0的基础组合
   • 后续块：依次针对每个浮点列，生成该列取非零值、其余浮点列取0的组合，再与布尔组合做笛卡尔积
3. 合并所有合法块，直接得到最终结果

优化代码 #

import pandas as pd

# 1. 生成布尔列的所有笛卡尔积组合
bool_combinations = pd.DataFrame([
    (a_val, b_val) for a_val in inputs["a"] for b_val in inputs["b"]
], columns=list(bool_inputs))

# 2. 构建所有合法的浮点组合块
float_blocks = []
# 加入所有浮点列都为0的基础块
base_block = bool_combinations.copy()
for col in float_inputs:
    base_block[col] = 0.0
float_blocks.append(base_block)

# 针对每个浮点列，生成仅该列非零的组合块
for target_col in float_inputs:
    # 提取当前列的非零值
    non_zero_values = [val for val in inputs[target_col] if val != 0.0]
    for val in non_zero_values:
        block = bool_combinations.copy()
        # 设置当前列值，其余浮点列置0
        for col in float_inputs:
            block[col] = val if col == target_col else 0.0
        float_blocks.append(block)

# 3. 合并所有块并调整列顺序与输入一致
final_df = pd.concat(float_blocks, ignore_index=True)
final_df = final_df[list(inputs.keys())]

方案优势 #

• 内存高效：直接生成合法行，避免全量生成后过滤带来的内存浪费，内存占用仅为合法行的大小
• 类型正确：布尔列类型天然匹配，无需额外转换
• 逻辑清晰：分块生成符合约束的组合，避免无效计算

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