使用np.put处理多维NumPy数组时出现异常行为的原因及解决方法

为什么np.put会得到不符合预期的结果？ #

np.put的设计逻辑是针对数组扁平化后的一维索引进行赋值操作，它并不支持直接传入多维坐标对（如(0,1)这类行-列索引）作为参数。

当你传入多维坐标列表I = [(0, 1), (1, 0), (1, 1)]时，numpy会先将这个列表转换为二维数组，再自动扁平化处理，得到一维索引序列[0,1,1,0,1,1]。随后np.put会循环使用v = [1,2,3]的值对这些一维索引位置赋值：

• 索引0（对应原数组的(0,0)位置）被赋值1
• 索引1（对应原数组的(0,1)位置）先后被赋值2、3，最终保留3
• 其他未被覆盖的位置保持初始值0

这就导致你得到了[[1., 3.], [0., 0.]]的结果，完全偏离了预期的多维坐标赋值逻辑。

如何实现预期的多维索引赋值？ #

方法一：直接使用多维索引赋值 #

将坐标列表拆分为行索引数组和列索引数组，直接通过numpy的多维索引语法赋值：

import numpy as np

A = np.array([[0, 0], [0, 0]])
I = [(0, 1), (1, 0), (1, 1)]
v = [1, 2, 3]

# 拆分坐标为行、列索引
rows, cols = zip(*I)
A[rows, cols] = v

print(A)

输出结果：

[[0 1]
 [2 3]]

方法二：转换为一维索引后使用np.put #

如果一定要使用np.put，可以先通过np.ravel_multi_index将多维坐标转换为扁平化的一维索引，再传入np.put：

import numpy as np

A = np.array([[0, 0], [0, 0]])
I = [(0, 1), (1, 0), (1, 1)]
v = [1, 2, 3]

# 将多维坐标转换为扁平化一维索引
flat_indices = np.ravel_multi_index(np.array(I).T, A.shape)
np.put(A, flat_indices, v)

print(A)

输出结果同样符合预期：

[[0 1]
 [2 3]]

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