如何高效生成固定长度N且含M个1的0-1数组所有组合？

纯Python简洁实现（基于itertools） #

直接利用标准库itertools.combinations生成所有M个索引的组合，再构造对应的0-1字符串或数组，底层是C实现，比手动多层循环效率高得多：

import itertools

def generate_01_combinations(N, M):
    for indices in itertools.combinations(range(N), M):
        # 构造0-1字符串
        combo = ['0'] * N
        for i in indices:
            combo[i] = '1'
        yield ''.join(combo)

# 测试N=4, M=2的场景
for idx, s in enumerate(generate_01_combinations(4, 2), 1):
    print(f"{idx}: '{s}'")

结合Numpy的矢量化实现（适合大规模场景） #

如果需要生成Numpy数组格式的结果，可通过批量初始化+索引赋值的方式优化，减少Python层循环开销：

import itertools
import numpy as np

def generate_numpy_01_arrays(N, M):
    total = np.math.comb(N, M)
    # 初始化全0的二维数组，行数为组合总数，列数为N
    result_arr = np.zeros((total, N), dtype=int)
    for idx, positions in enumerate(itertools.combinations(range(N), M)):
        result_arr[idx, positions] = 1
    return result_arr

# 测试并输出字符串格式结果
arrays = generate_numpy_01_arrays(4, 2)
for idx, row in enumerate(arrays, 1):
    print(f"{idx}: '{''.join(map(str, row))}'")

关键优势 #

• itertools.combinations是标准库内置工具，无需额外依赖，执行效率远高于手动嵌套循环。
• Numpy方案通过矢量化操作处理批量数据，在N和M取值较大时，性能提升更显著。

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