如何使用argwhere函数按列访问二进制numpy数组的索引

现有如下Numpy二进制数组：

import numpy as np
x=np.array([[1, 1, 0, 0, 1, 0, 0, 0, 0, 0],
    [0, 0, 0, 1, 0, 1, 0, 0, 0, 0],
    [0, 0, 1, 1, 0, 0, 1, 0, 0, 0],
    [0, 1, 0, 1, 0, 0, 0, 0, 1, 0],
    [1, 0, 0, 0, 0, 0, 1, 0, 0, 0]])        

使用ind=np.argwhere(x)得到的是按行排列的索引，现在需要获取按列排列的非零元素索引，期望输出如下：

[[0 0]
 [4 0]
 [0 1]
 [3 1]
 [2 2]
 [1 3]
 [2 3]
 [3 3]
 [0 4]
 [1 5]
 [2 6]
 [4 6]
 [3 8]]

解决方案 #

可以通过对argwhere的结果按列索引排序实现，具体代码如下：

import numpy as np

x = np.array([[1, 1, 0, 0, 1, 0, 0, 0, 0, 0],
              [0, 0, 0, 1, 0, 1, 0, 0, 0, 0],
              [0, 0, 1, 1, 0, 0, 1, 0, 0, 0],
              [0, 1, 0, 1, 0, 0, 0, 0, 1, 0],
              [1, 0, 0, 0, 0, 0, 1, 0, 0, 0]])        

# 获取所有非零元素的索引
ind = np.argwhere(x)
# 按列索引（数组第二列，索引为1）排序
sorted_ind = ind[ind[:, 1].argsort()]

print(sorted_ind)

原理说明 #

• np.argwhere(x)会返回所有非零元素的[行索引, 列索引]对，默认按行优先顺序排列。
• ind[:, 1]提取出所有列索引，argsort()返回这些列索引升序排列对应的位置下标。
• 用该下标对ind数组切片，即可得到按列索引排序的结果，同一列内的行索引会保留原顺序，与期望输出完全匹配。

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