如何用Numpy高效创建斜向填充的升序整数二维矩阵？

矢量化解决方案 #

完全基于NumPy的矢量化操作实现，避免Python循环开销，大幅提升大N场景下的性能：

import numpy as np

def create_triangular_matrix(N):
    # 计算矩阵边长m：满足m(m+1)/2 >= N的最小整数
    m = np.ceil((np.sqrt(8 * N + 1) - 1) / 2).astype(int)
    # 生成1到N的整数数组
    ks = np.arange(1, N + 1)
    # 计算每个数所在的斜线s（s = x + y，斜线的对角线和）
    s = np.floor((np.sqrt(8 * ks - 7) - 1) / 2).astype(int)
    # 计算每条斜线的起始数值
    start = s * (s + 1) // 2
    # 计算当前数在斜线中的偏移位置
    offset = ks - start - 1
    # 推导每个数的矩阵坐标
    y = s - offset
    x = offset
    # 创建零矩阵并批量赋值
    mat = np.zeros((m, m), dtype=int)
    mat[y, x] = ks
    return mat

# 测试示例
print("N=6时的矩阵：")
print(create_triangular_matrix(6))
print("\nN=13时的矩阵：")
print(create_triangular_matrix(13))

核心逻辑说明 #

1. 矩阵边长计算：通过二次方程求解，得到能容纳所有1~N数值的最小正方形矩阵边长m，确保空间足够且无冗余。
2. 斜线分组：每个数值k对应一条对角线和为s的斜线，斜线s包含s+1个连续数值，起始值为s*(s+1)//2 + 1。
3. 坐标推导：每条斜线的数值从矩阵左下角（y=s, x=0）开始，依次向上右方向填充，通过偏移量计算每个k的精确坐标。
4. 批量赋值：利用NumPy的索引广播特性，一次性完成所有位置的赋值，彻底避免循环带来的性能损耗。

性能优势 #

对于大N（如N=10^6），循环版本会因Python解释器的开销变得极慢，而矢量化方案依赖NumPy底层C实现的操作，能在毫秒级完成计算，性能提升可达数百倍。

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