Python中如何仅用单循环或切片遍历numpy二维矩阵并完成条件赋值

你提到的场景完全可以摆脱Python层的双层嵌套循环实现，既可以做到零循环的原生切片操作（性能远优于逐元素遍历），也可以用单循环适配更复杂的多规则赋值逻辑。 #

一、零循环切片实现（针对固定偏移对角线的赋值场景） #

最初你提到的给主对角线赋值4的场景，仅需要调用一次np.fill_diagonal(A, 4)即可完成，完全不需要循环。而你补充的三类赋值规则本质都是矩阵固定偏移量的对角线元素赋值，同样可以通过NumPy内置的对角线定位能力直接完成，不需要任何循环：

• 主对角线（i == j）：矩阵中从左上角延伸到右下角的核心对角线
• 上邻次对角线（i +1 == j）：主对角线向右侧偏移1位的对角线
• 下邻次对角线（i -1 == j）：主对角线向左侧偏移1位的对角线

对应的实现代码和你的双层循环输出完全等价：

import numpy as np
A = np.zeros((10, 10))

# 主对角线赋值为1
np.fill_diagonal(A, 1)
# 偏移量为1的上邻次对角线赋值为2
np.fill_diagonal(A, 2, offset=1)
# 偏移量为-1的下邻次对角线赋值为3
np.fill_diagonal(A, 3, offset=-1)

如果不想依赖fill_diagonal方法，也可以通过广播索引手动定位对应元素，效果完全一致：

# 主对角线定位
A[np.arange(10), np.arange(10)] = 1
# 上邻次对角线定位（i最大为8，避免j=10越界）
A[np.arange(9), np.arange(9)+1] = 2
# 下邻次对角线定位（i最小为1，避免j=-1越界）
A[np.arange(1, 10), np.arange(9)] = 3

这两类实现的底层都是NumPy在C层直接定位目标元素，没有Python层的逐元素遍历，执行效率比双层Python循环高1~2个数量级。 #

二、单循环广播实现（针对通用多规则赋值场景） #

如果你的赋值规则不是固定偏移的对角线，而是任意组合的条件判断，也可以用单循环完全替代双层嵌套遍历：
核心思路是一次性预生成全矩阵的行、列索引矩阵，之后把每个赋值规则作为布尔判断条件，直接通过NumPy的布尔索引批量给符合条件的位置赋值，整个循环的次数仅等于赋值规则的数量，和矩阵尺寸无关。
对应你提到的多规则场景实现：

A = np.zeros((10, 10))
# 一次性生成全矩阵的行索引i、列索引j，尺寸和A完全一致
i, j = np.mgrid[0:10, 0:10]

# 单循环遍历所有赋值规则，无嵌套
rules = [
    (i == j, 1),
    (i + 1 == j, 2),
    (i - 1 == j, 3)
]
for condition, value in rules:
    A[condition] = value

这个实现和你最初的双层循环输出完全一致，但完全避免了原来双层循环逐元素判断的O(n²)Python层开销，哪怕是尺寸达到1000×1000的矩阵，整个循环也只需要3次迭代。 #

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