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

【问题标题】：numpy most efficient way to get row index of true values by columnsnumpy 按列获取真值行索引的最有效方法
【发布时间】：2017-03-23 09:06:17
【问题描述】：

我想从二维 ndarray 中按列获取真值的行索引。到目前为止，我有一个带有 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时，我需要得到一个空列表[]而不是跳过。

【问题讨论】：

这看起来很像 scipy.sparse lil 格式。

标签： python numpy

【解决方案1】：

方法#1

这是一种矢量化 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]]

方法 #2

或者，我们可以使用loop comprehension 来避免这种分裂 -

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 [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

这是处理所有0s cols 的另一种方法-

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。

【讨论】：

你的解决方案有点缺陷，当一列都是False时，我需要得到空列表[]。
避免拆分使您的解决方案非常快！
非常感谢！方法＃3正是我想要的。顺便说一句，你的智慧令人难以置信！我还回答了带有numpy 和pandas 标记的问题，并看到了许多您的精彩答案。我真的很想知道你是如何学习这两个库的。
我会将r,c = np.where(arr.T) 更改为c, r = np.where(arr)。不能保证转置给出一个视图。所以在某些情况下，第二种变体可能会更快。
@gzc 我从 MATLAB 开始。然后，NumPy 是一个简单的扩展，然后 pandas 在某种程度上处理 NumPy，所以也开始了，通过尝试在这里回答问题 :)

【解决方案2】：

我的解决方案是

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

正如所指出的，我们仍然缺少错误的列，这些可以使用 unique 的信息插入：

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 部分，因此会立即清楚发生了什么。

【讨论】：

请检查我的笔记。 Divakar 的方法#1 和#2 是有缺陷的。你的也是。