Numpy中如何生成给定形状（如(10,3)）的所有可广播形状？

解决方案 #

要生成所有能与给定形状广播的至少1维形状，我们可以利用NumPy的矢量化操作来高效实现，核心思路是紧扣广播规则推导符合条件的形状特征：

1. 广播规则核心：一个k维形状能和n维原形状广播，等价于在k维形状左侧补(n-k)个1后，得到的n维形状每个维度要么等于原形状对应维度，要么是1。
2. 生成基础n维形状：先生成所有n维候选形状，每个维度可选1或原形状的对应维度（共2ⁿ种组合）。
3. 提取有效后缀：从每个基础n维形状中提取所有非空后缀（去掉前面的若干个1），这些后缀就是符合要求的k维形状。
4. 去重与排序：去除重复形状后，按维度长度、形状内容排序得到最终结果。

代码实现

import numpy as np

def get_broadcastable_shapes(original_shape):
    dims = np.array(original_shape)
    n = len(dims)
    
    # 为每个维度生成两种选择：1 或 原维度值
    choices = np.stack([np.ones(n, dtype=int), dims], axis=1)
    
    # 生成所有维度选择的笛卡尔积索引（每个维度选0对应1，选1对应原维度）
    indices = np.indices([2]*n).reshape(n, -1).T
    
    # 批量生成所有符合条件的n维形状
    n_dimensional_shapes = np.take_along_axis(choices, indices, axis=1)
    
    # 收集所有非空后缀形状
    all_suffixes = []
    for start in range(n):
        # 提取每个n维形状从start位置开始的后缀，转为元组便于去重
        suffixes = [tuple(shape) for shape in n_dimensional_shapes[:, start:]]
        all_suffixes.extend(suffixes)
    
    # 去重并排序：先按维度长度排序，长度相同则按形状内容排序
    unique_shapes = sorted(list(set(all_suffixes)), key=lambda x: (len(x), x))
    
    return unique_shapes

# 测试示例
original_shape = (10, 3)
print(get_broadcastable_shapes(original_shape))
# 输出：[(1,), (3,), (1, 1), (1, 3), (10, 1), (10, 3)]

代码说明

• 矢量化生成基础形状：通过np.indices和np.take_along_axis批量生成所有n维候选形状，这一步依赖NumPy的C后端实现，比纯Python循环生成笛卡尔积效率高得多。
• 提取后缀覆盖所有维度：遍历每个可能的起始位置，确保提取出1到n维的所有有效形状。
• 去重排序保证结果整洁：用集合快速去重，再通过自定义排序规则让结果更易读。

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