如何在Numpy中递归生成尺寸逐次翻倍的极化码生成矩阵

核心实现思路 #

你需要的矩阵是极化码的生成矩阵，构造规则对应分块矩阵拼接逻辑，完全可以用Numpy内置的向量化操作实现，不需要手写元素级的双层循环，效率更高也更简洁。

实现方案1：分块直接拼接 #

直接用np.block对上层矩阵、零矩阵做分块拼接，完全匹配你给出的构造规则：

import numpy as np

def generate_polar_matrix(k: int):
    """
    生成尺寸为 2^k * 2^k 的极化码生成矩阵
    参数k: 迭代次数，k=1返回初始2阶矩阵，k=2返回4阶矩阵，以此类推
    """
    # 初始2阶生成矩阵
    g = np.array([[1, 0], [1, 1]], dtype=np.int8)
    if k == 1:
        return g
    for _ in range(k - 1):
        zero_mat = np.zeros_like(g)
        # 按规则拼接分块
        g = np.block([
            [g, zero_mat],
            [g, g]
        ])
    return g

调用示例：生成4阶矩阵直接调用generate_polar_matrix(2)，返回结果和你给出的示例完全一致。

实现方案2：克罗内克积实现（更简洁） #

极化码生成矩阵本质是基础2阶矩阵的n次克罗内克幂，用np.kron实现代码更短：

import numpy as np

def generate_polar_matrix_kronecker(k: int):
    base = np.array([[1, 0], [1, 1]], dtype=np.int8)
    res = base.copy()
    for _ in range(k - 1):
        res = np.kron(res, base)
    return res

优化建议 #

• 如果只需要0/1的布尔值标识，可以将dtype设置为np.bool_，内存占用可以降低到原来的1/8，大尺寸矩阵下性能提升明显。
• 两种实现均为向量化操作，相比手写双层for循环填充，在尺寸大于1024*1024的场景下，性能提升可达百倍以上。

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