如何使用Python将Numpy数组的下三角（含对角线）填充为-inf？

核心思路 #

利用Numpy内置的下三角索引生成函数，直接定位到需要修改的区域后批量赋值即可，时间复杂度为O(n)，操作效率极高。

完整实现代码 #

import numpy as np

# 1. 构造原始数组，注意：若原始数组为整数类型，需先转为浮点型（int类型不支持inf存储）
arr = np.array([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
], dtype=np.float64)

# 2. 获取下三角（包含对角线）的坐标索引
rows, cols = np.tril_indices(arr.shape[0])

# 3. 对应位置赋值为负无穷
arr[rows, cols] = -np.inf

# 输出验证结果
print(arr)

运行输出 #

[[-inf   2.   3.]
 [-inf -inf   6.]
 [-inf -inf -inf]]

拓展说明 #

• 若处理的是行列数不等的2D数组，只需要在tril_indices中传入列数参数即可：

  rows, cols = np.tril_indices(arr.shape[0], n=arr.shape[1])
  

• 也可以用布尔掩码的方式实现同样效果，代码更简洁：

  arr[np.tri(*arr.shape, dtype=bool)] = -np.inf
  

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