使用 numpy 获取沿另一列排序的给定列值的行索引答案

【问题标题】：Use numpy to get row indexes for a given column value sorted along another column使用 numpy 获取沿另一列排序的给定列值的行索引
【发布时间】：2019-12-02 18:32:19
【问题描述】：

问题的标题可能令人困惑，但问题出在这里，我有二维 numpy 数组。现在，我想获取第一列具有特定值的索引列表/数组，同时沿第二列排序：

a = np.array([[1,2],[1,3],[1,4],[1,5],[1,6],[2,9],[1,9],[1,7],[2,7],[1,8]])

index = [0, 1, 2, 3, 4, 7, 9, 6] # <---- the solution, I want this list

# this list gives sorted array for 1st column value 1
a[index] = 
array([[1, 2],
       [1, 3],
       [1, 4],
       [1, 5],
       [1, 6],
       [1, 7],
       [1, 8],
       [1, 9]])

注意：我想要索引列表，而不是给定值的排序数组。

我目前想出的如下：

tmp = a[np.lexsort((a[:,1],a[:,0]))]
tmp= tmp[tmp[:,0]==1]
index = [np.where(np.all(a==i,axis=1))[0][0] for i in tmp]

你可以看到这很糟糕，而且我正在处理非常大的数据集，这需要真正的改进。有没有什么方法可以用 numpy 更有效地完成这个任务？

【问题讨论】：

如果您的问题得到解决，您应该标记正确答案。
发布的解决方案是否对您有用？对已发布的解决方案有何反馈？

标签： python arrays numpy

【解决方案1】：

使用np.lexsort -

# Perform lex-sorting keeping the second col as the primary order.
# Thus, identical elements along that col will be in sequence.
# This helps in keeping unique ones in it later on w/ consecutive checks
In [355]: idx = np.lexsort(a.T)

# Index into the first col with idx, check for the identifier=1
# and then filter out those off idx
In [160]: out = idx[a[idx,0]==1]

In [161]: out
Out[161]: array([0, 1, 2, 3, 4, 7, 9, 6])

大型随机数组的计时 -

In [167]: np.random.seed(0)    
     ...: a = np.random.randint(0,1000,(100000,2))
     ...: a = np.unique(a,axis=0)
     ...: np.random.shuffle(a)

# @Akaisteph7's soln
In [168]: %%timeit
     ...: tmp = np.unique(a, return_index=True, axis=0)
     ...: index = tmp[1][tmp[0][:,0]==1]
10 loops, best of 3: 48.4 ms per loop

# From this post
In [169]: %%timeit
     ...: idx = np.lexsort(a.T)
     ...: out = idx[a[idx,0]==1]
10 loops, best of 3: 22.1 ms per loop

【讨论】：

能否详细说明。您将列索引放在哪里，值在哪里？
@Eular 认为我之前误解了。已编辑帖子。此外，您已经声明您正在处理大数据。您能否测试一下发布的方法在您的最后是否公平？

【解决方案2】：

这是使用np.unique 的另一种方法。 np.unique 的优点是您可以将其配置为直接返回索引和排序数组。见以下代码：

# Get the sorted array and indices
tmp = np.unique(a, return_index=True, axis=0)
# Get the indices only where the sorted array's first column equals 1 
index = tmp[1][tmp[0][:,0]==1]
print(index)

输出：

[0 1 2 3 4 7 9 6]

【讨论】：