Python 如何在numpy数组的每一行中将第n个非零元素更改为零

发布于06月26日

我有一个宽的二进制二维numpy数组，如下所示:

np_var:
0, 0, 1, 0, 1, ..., 0, 1
1, 0, 1, 0, 0, ..., 1, 0
...

每行有8个非零元素.我想用零快速替换每行中的第n个非零元素(最终每行有7个非零元素).

有没有一种简单的方法可以在没有循环的情况下快速执行此替换？

推荐答案

可以获取非零元素的索引，并使用它们替换数组中的值

arr = np.array(...)
print(arr)

# [[1 1 1 0 0 1 0 1 1 0 0 1 1 0]
#  [0 1 1 1 1 0 1 0 0 1 1 0 1 0]
#  [1 0 1 1 0 1 1 1 0 1 0 0 1 0]
#  [0 1 1 0 1 0 0 1 1 1 1 0 1 0]
#  [1 1 1 0 0 1 1 0 0 1 1 0 0 1]
#  [0 0 1 1 1 1 1 0 1 1 0 0 1 0]
#  [1 0 1 0 1 0 1 1 1 0 0 1 0 1]
#  [1 0 1 1 1 0 1 1 0 0 1 0 1 0]
#  [0 0 1 1 1 1 0 1 0 1 1 0 0 1]
#  [0 1 1 1 0 0 0 1 1 0 1 1 1 0]]

nth_element = 5
non_zero_count = int(np.count_nonzero(arr) / len(arr)) # can be replaced by 8 if the size is fixed
indices = arr.nonzero()[1][nth_element - 1::non_zero_count]
arr[np.arange(len(arr)), indices] = 5
print(arr)

# [[1 1 1 0 0 1 0 5 1 0 0 1 1 0]
#  [0 1 1 1 1 0 5 0 0 1 1 0 1 0]
#  [1 0 1 1 0 1 5 1 0 1 0 0 1 0]
#  [0 1 1 0 1 0 0 1 5 1 1 0 1 0]
#  [1 1 1 0 0 1 5 0 0 1 1 0 0 1]
#  [0 0 1 1 1 1 5 0 1 1 0 0 1 0]
#  [1 0 1 0 1 0 1 5 1 0 0 1 0 1]
#  [1 0 1 1 1 0 5 1 0 0 1 0 1 0]
#  [0 0 1 1 1 1 0 5 0 1 1 0 0 1]
#  [0 1 1 1 0 0 0 1 5 0 1 1 1 0]]