Python中指定行置零：numpy数组行处理报错解决方案

我有一个数组列表Pe，想要保留列表J中指定索引的行，其余行全部置零。例如Pe[0]对应J[0]=[0,1]，需保留Pe[0]的第0、1行，第2行全置零；Pe[1]同理。但运行代码时出现错误，以下是相关内容：

原代码 #

import numpy as np

Pe = [np.array([[402.93473651,   0.        , 230.97804127, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        , 423.81345923,   0.        , 407.01354328,
        419.14952534,   0.        , 316.58460442,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [402.93473651,   0.        , 230.97804127, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ]]),
np.array([[402.93473651,   0.        , 230.97804127, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        , 423.81345923,   0.        , 407.01354328,
        419.14952534,   0.        , 316.58460442,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [402.93473651,   0.        , 230.97804127, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ]])]  #Entry pressure

J = [[0,1],[2]]

for i in range(0,len(Pe)):
    out = np.zeros_like(Pe[i])
    for j in range(0,len(J)):
        out[i][J[j]] = Pe[i][J[j]]
    print([out])

错误信息 #

in <module>
    out[i][J[j]] = Pe[i][J[j]]

ValueError: shape mismatch: value array of shape (2,12)  could not be broadcast to indexing result of shape (2,)

预期输出 #

[np.array([[402.93473651,   0.        , 230.97804127, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        , 423.81345923,   0.        , 407.01354328,
        419.14952534,   0.        , 316.58460442,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [0. ,            0.        ,   0. ,          0. ,
          0.        , 0. ,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ]]),
np.array([[0. ,   0.        , 0. , 0. ,
          0.        , 0. ,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        , 0. ,   0.        , 0. ,
        0. ,   0.        , 0.,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [402.93473651,   0.        , 230.97804127, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ]])]

错误原因 #

代码中嵌套循环的索引逻辑错误：

• 外层循环的i是遍历Pe的每个二维数组，但out[i][J[j]]是取当前零数组的第i行，再选中该行的J[j]列；而Pe[i][J[j]]是取原数组的J[j]行，两者形状（前者是一维数组，后者是二维数组）不匹配，导致报错。

修正后的代码 #

import numpy as np

Pe = [np.array([[402.93473651,   0.        , 230.97804127, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        , 423.81345923,   0.        , 407.01354328,
        419.14952534,   0.        , 316.58460442,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [402.93473651,   0.        , 230.97804127, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ]]),
np.array([[402.93473651,   0.        , 230.97804127, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        , 423.81345923,   0.        , 407.01354328,
        419.14952534,   0.        , 316.58460442,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [402.93473651,   0.        , 230.97804127, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ]])]

J = [[0,1],[2]]

result = []
for arr, keep_indices in zip(Pe, J):
    out = np.zeros_like(arr)
    out[keep_indices] = arr[keep_indices]
    result.append(out)

print(result)

说明 #

• 用zip(Pe, J)直接配对每个数组和对应的保留索引列表，避免嵌套循环的索引混乱
• out[keep_indices]直接选中要保留的整行，赋值为原数组的对应行，形状完全匹配，不会触发报错
• 最终将处理后的数组收集到result列表中，输出结果与预期一致

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