如何高效生成含n个1且行列无重复1的p×q矩阵及计数公式

需求：生成所有p×q矩阵，要求矩阵恰好包含n个"1"，且每行、每列最多存在一个"1"（类似国际象棋n车问题）。以下是p=4、q=4、n=2时的低效Python实现：

for i1 in range(4*4):
    x1, y1 = i1//4, i1%4
    for i2 in range(i1+1,4*4):  # 注：原代码中i应为i1，此处修正笔误
        x2, y2 = i2//4, i2%4
        if x1 != x2 and y1 != y2:
           print(x1,y1,x2,y2)

该方法存在效率低、n增大时逻辑繁琐的问题，疑问如下：

• 是否存在向量化的实现方式？
• 这类矩阵的计数公式是什么？

一、高效实现方案（含向量化优化） #

可以利用组合排列的特性直接生成所有合法的"1"位置，避免无效遍历和判断。以下是基于itertools和numpy的优化实现：

import itertools
import numpy as np

p, q, n = 4, 4, 2

# 生成所有合法的"1"坐标组合
# 1. 从p行中选n个不重复的行
row_combinations = itertools.combinations(range(p), n)
# 2. 对每个行组合，生成n个不重复列的全排列（保证列唯一）
valid_positions = []
for rows in row_combinations:
    for cols in itertools.permutations(range(q), n):
        valid_positions.append(list(zip(rows, cols)))

# 向量化生成矩阵
matrices = []
for pos in valid_positions:
    mat = np.zeros((p, q), dtype=int)
    # 用numpy向量化赋值替代循环
    coords = np.array(pos).T
    mat[coords[0], coords[1]] = 1
    matrices.append(mat)

# 验证输出（打印前3个矩阵）
for idx, mat in enumerate(matrices[:3], 1):
    print(f"矩阵{idx}:")
    print(mat)

优势说明 #

• 无需嵌套多层循环，n值变化时仅需调整参数，逻辑简洁
• 直接通过组合排列筛选合法位置，避免原方案中大量无效遍历和条件判断
• 利用numpy向量化赋值进一步提升矩阵生成效率

二、计数公式 #

这类矩阵的总数可通过组合数学推导得出：

1. 从p行中挑选n个不同行的组合数为 $\binom{p}{n} = \frac{p!}{n!(p-n)!}$
2. 给选中的n个行分配n个不同的列，等价于从q列中选n个列的排列数 $P(q, n) = \frac{q!}{(q-n)!}$

两者相乘即为总数量：
$$
\text{总数} = \binom{p}{n} \times P(q, n) = \frac{p! \times q!}{n!(p-n)!(q-n)!}
$$

也可等价表示为从q列选n个组合，再给n个行做排列：
$$
\text{总数} = P(p, n) \times \binom{q}{n} = \frac{p! \times q!}{(p-n)!n!(q-n)!}
$$

示例验证 #

当p=4、q=4、n=2时：
$\binom{4}{2} = 6$，$P(4,2)=12$，总数为 $6 \times 12 = 72$，与原方案遍历去重后的结果一致。

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