基于NumPy矩阵生成自定义输出的Python实现方案求助

首先，我们可以利用NumPy的向量化操作高效实现所有i值的需求，核心思路是先提取上三角（主对角线以上）的元素，再根据i值筛选需要保留的元素，其余置0：

import numpy as np

# 生成原始矩阵M
M = np.arange(25).reshape(-1, 5) * 10

def get_custom_list(i):
    # 获取上三角（k=1，即主对角线右侧）的行、列索引
    r, c = np.triu_indices(5, k=1)
    # 生成掩码：保留列索引 <= i-1 的元素
    mask = c <= (i - 1)
    # 提取上三角元素，将掩码为False的位置置0，转成列表
    upper_tri_vals = M[r, c]
    upper_tri_vals[~mask] = 0
    return upper_tri_vals.tolist()

# 验证各i值的输出
print(get_custom_list(5))  # [10, 20, 30, 40, 70, 80, 90, 130, 140, 190]
print(get_custom_list(4))  # [10, 20, 30, 0, 70, 80, 0, 130, 0, 0]
print(get_custom_list(3))  # [10, 20, 0, 0, 70, 0, 0, 0, 0, 0]
print(get_custom_list(2))  # [10, 0, 0, 0, 0, 0, 0, 0, 0, 0]

思路解释 #

1. 矩阵结构分析：原始矩阵M是5×5的矩阵，我们需要的元素都位于主对角线的右侧（上三角区域，k=1）。
2. 索引提取：用np.triu_indices(5, k=1)直接获取上三角所有元素的行、列索引，避免手动遍历。
3. 掩码筛选：观察各i值的输出规律，发现保留的元素满足列索引 ≤ i-1，通过这个条件生成布尔掩码，快速标记需要保留的位置。
4. 置0与转换：提取上三角元素后，将掩码为False的位置置0，最后转换为Python列表，完成需求。

这种方式完全基于NumPy的向量化操作，避免了循环，效率远高于逐元素判断的实现。

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