NumPy实现N边长0菱形全1矩阵 寻求无循环向量化优化方案

需求说明 #

• 入参为大于0的整数N，返回尺寸为(2N-1)×(2N-1)的二维数组
• 数组默认元素全为1，其中值为0的元素组成菱形图案，菱形4条边的单条长度均为N
• 实现约束：禁止使用循环、列表推导、递归，需充分利用NumPy向量化特性，避免冗余的分块拼接逻辑

预期运行结果 #

• 当N=1时返回：

np.array([[0]])

• 当N=2时返回：

np.array([[1,0,1],[0,1,0],[1,0,1]])

• 当N=3时返回：

np.array([[1,1,0,1,1],[1,0,1,0,1],[0,1,1,1,0],[1,0,1,0,1],[1,1,0,1,1]])

原有实现问题 #

原有方案通过全1矩阵减对角线、四分块镜像拼接的方式构造结果，逻辑偏绕，还单独对N=1做了边界判断，本质是手动拆分矩阵区域赋值，没有发挥NumPy批量坐标计算的优势。

向量化实现方案 #

菱形四条边刚好对应两类对角线位置：一类是行列索引差等于中心偏移量，另一类是行列索引和与中心行列和的差等于中心偏移量。直接通过坐标网格批量标记0值位置即可，无需分块，也不需要单独处理边界：

import numpy as np

def generate_diamond(n: int) -> np.ndarray:
    size = 2 * n - 1
    center = n - 1
    # 生成全矩阵的行、列坐标网格
    row_idx, col_idx = np.mgrid[:size, :size]
    # 标记所有菱形边的位置（即需要置0的位置）
    zero_mask = (np.abs(row_idx - col_idx) == center) | (np.abs(row_idx + col_idx - 2 * center) == center)
    # 初始化全1矩阵，标记位置置0
    result = np.ones((size, size), dtype=np.int64)
    result[zero_mask] = 0
    return result

这个实现全程使用NumPy内置的向量化操作，没有任何循环、递归、列表推导逻辑，所有N≥1的场景都能直接返回正确结果，代码逻辑清晰，执行效率更高。

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