Python numpy按列获取真值行索引的最有效方法

我想从二维数组中按列获取真值的行索引。到目前为止，我有一个带有for循环的解决方案。但我认为这是没有效率的，因为存在一个python本机for循环。我试图找出一个矢量化的解决方案，但失败了

更新：无需是矢量化解决方案，效率越高越好

arr = np.random.randint(2, size=15).reshape((3,5)).astype(bool)
print arr

[[ True False  True False  True]
 [False  True False  True  True]
 [ True  True False False  True]]

def calc(matrix):
    result = []
    for i in range(matrix.shape[1]):
        result.append(np.argwhere(matrix[:, i]).flatten().tolist())
    return result

print calc(arr)
[[0, 2], [1, 2], [0], [1], [0, 1, 2]]

注意：我希望行索引按列分组。当一列全部为False时，我需要得到一个空列表，而不是跳过

这里有一种矢量化NumPy方法，可以将这些行索引分组到数组列表中-

r,c = np.where(arr.T)
out = np.split(c, np.flatnonzero(r[1:] != r[:-1])+1)

样本运行-

In [63]: arr = np.random.randint(2, size=15).reshape((3,5)).astype(bool)

In [64]: arr
Out[64]: 
array([[False, False,  True,  True, False],
       [ True,  True, False, False,  True],
       [ True,  True, False, False,  True]], dtype=bool)

In [65]: r,c = np.where(arr.T)

In [66]: np.split(c, np.flatnonzero(r[1:] != r[:-1])+1)
Out[66]: [array([1, 2]), array([1, 2]), array([0]), array([0]), array([1, 2])]

In [67]: calc(arr)
Out[67]: [[1, 2], [1, 2], [0], [0], [1, 2]]

In [177]: arr
Out[177]: 
array([[ True, False, False, False, False],
       [ True, False, False, False,  True],
       [ True, False,  True, False,  True]], dtype=bool)

In [178]: idx = np.concatenate(([0], arr.sum(0).cumsum() ))
     ...: out = [c[idx[i]:idx[i+1]] for i in range(len(idx)-1)]
     ...: 

In [179]: out
Out[179]: 
[array([0, 1, 2]),
 array([], dtype=int64),
 array([2]),
 array([], dtype=int64),
 array([1, 2])]

方法#2

或者，我们可以使用

循环理解
来避免这种分裂-

idx = np.concatenate(([0], np.flatnonzero(r[1:] != r[:-1])+1, [r.size] ))
out = [c[idx[i]:idx[i+1]] for i in range(len(idx)-1)]

我们使用方法1中的
r，c

方法#3（为所有0列输出空列表/数组）

为了解释所有的零列，我们需要空列表/数组，这里有一种改进的方法-

idx = np.concatenate(([0], arr.sum(0).cumsum() ))
out = [c[idx[i]:idx[i+1]] for i in range(len(idx)-1)]

我们使用的是方法1中的
c

样本运行-

In [63]: arr = np.random.randint(2, size=15).reshape((3,5)).astype(bool)

In [64]: arr
Out[64]: 
array([[False, False,  True,  True, False],
       [ True,  True, False, False,  True],
       [ True,  True, False, False,  True]], dtype=bool)

In [65]: r,c = np.where(arr.T)

In [66]: np.split(c, np.flatnonzero(r[1:] != r[:-1])+1)
Out[66]: [array([1, 2]), array([1, 2]), array([0]), array([0]), array([1, 2])]

In [67]: calc(arr)
Out[67]: [[1, 2], [1, 2], [0], [0], [1, 2]]

In [177]: arr
Out[177]: 
array([[ True, False, False, False, False],
       [ True, False, False, False,  True],
       [ True, False,  True, False,  True]], dtype=bool)

In [178]: idx = np.concatenate(([0], arr.sum(0).cumsum() ))
     ...: out = [c[idx[i]:idx[i+1]] for i in range(len(idx)-1)]
     ...: 

In [179]: out
Out[179]: 
[array([0, 1, 2]),
 array([], dtype=int64),
 array([2]),
 array([], dtype=int64),
 array([1, 2])]

方法#4

下面是处理所有
0
col的另一种方法-

unq, IDs = np.unique(r, return_index=1)
idx = np.concatenate(( IDs, [r.size] ))
out = [[]]*arr.shape[1]
for i,item in enumerate(unq):
    out[item] = c[idx[i]:idx[i+1]]

我们正在使用方法1中的
r，c
。
我的解决方案是

column, row = np.where(arr.T)
unique, indices = np.unique(column, return_index=True)
res = np.split(row, indices[1:])

如前所述，我们仍然缺少虚假的栏目，
可以使用以下信息插入这些文件：

missing = np.setdiff1d(np.arange(arr.shape[-1]), unique)
for mm in missing:
    res.insert(mm, np.array([], dtype=int))
   

这比@Divakar提议的要慢一点。但是，我发现它更具可读性，因为复杂的
np.flatnonzero（r[1::！=r[：-1]）+1
部分可以避免，所以发生的事情马上就清楚了。
而且您的解决方案有点缺陷，当一列都是
False
时，我需要得到空列表
[]
。避免这种拆分会让您的解决方案非常快！非常感谢你！方法3正是我想要的。顺便说一句，你的智慧是难以置信的！我还回答了标有
numpy
和
pandas
的问题，并看到了您的许多精彩答案。我真想知道你是如何学习这两个库的。我会把
r，c=np.where（arr.T）
改为
c，r=np.where（arr）
。不能保证转置可以提供视图。因此，在某些情况下，第二个变量可能要快得多。然后，NumPy是它的一个简单的扩展，然后熊猫在某种程度上与NumPy打交道，所以也开始了，通过尝试回答这里的问题：）请检查我的笔记。Divakar的方法1和方法2存在缺陷。你的也是。这看起来很像
scipy.sparse
lil
格式。