Question

我想根据第二列中的值仅选择NumPy数组中的某些行。例如，此测试数组在第二列中具有从1到10的整数。

>>> test = numpy.array([numpy.arange(100), numpy.random.randint(1, 11, 100)]).transpose()
>>> test[:10, :]
array([[ 0,  6],
       [ 1,  7],
       [ 2, 10],
       [ 3,  4],
       [ 4,  1],
       [ 5, 10],
       [ 6,  6],
       [ 7,  4],
       [ 8,  6],
       [ 9,  7]])

如果我只想要第二个值为4的行，那么很容易：

>>> test[test[:, 1] == 4]
array([[ 3,  4],
       [ 7,  4],
       [16,  4],
       ...
       [81,  4],
       [83,  4],
       [88,  4]])

但是，如果有多个想要的值，我如何获得相同的结果呢？

通缉名单可以是任意长度。例如，我可能想要第二列为2,4或6的所有行：

>>> wanted = [2, 4, 6]

我提出的唯一方法是使用列表理解然后将其转换回数组并且看起来过于复杂，尽管它有效：

>>> test[numpy.array([test[x, 1] in wanted for x in range(len(test))])]
array([[ 0,  6],
       [ 3,  4],
       [ 6,  6],
       ...
       [90,  2],
       [91,  6],
       [92,  2]])

在NumPy中有没有更好的方法可以做到这一点，我错过了？

Answer 1

随着wanted越来越大，以下解决方案应该比Amnon的解决方案更快：

# Much faster look up than with lists, for larger lists:
wanted_set = set(wanted)

@numpy.vectorize
def selected(elmt): return elmt in wanted_set
# Or: selected = numpy.vectorize(wanted_set.__contains__)

print test[selected(test[:, 1])]

事实上，它的优点是只搜索test数组一次（而不是像Amnon的答案那样多次len(wanted)次）。它还使用Python内置的快速元素查找 sets ，这比列表快得多。它也很快，因为它使用Numpy的快速循环。您还可以获得in运算符的优化：一旦wanted元素匹配，其余元素就不必进行测试（与Amnon的“逻辑或”方法相反，所有元素都是wanted无论如何都会被测试。

或者，您可以使用以下单行，也只能通过您的数组一次：

test[numpy.apply_along_axis(lambda x: x[1] in wanted, 1, test)]

这要慢很多，因为这会在每次迭代时提取第二列中的元素（而不是在一次传递中执行，就像在这个答案的第一个解决方案中一样）。

Answer 2

test[numpy.logical_or.reduce([test[:,1] == x for x in wanted])]

结果应该比原始版本更快，因为NumPy正在使用内部循环而不是Python。

Answer 3

numpy.in1d正是您所寻找的：

print test[numpy.in1d(test[:,1], wanted)]

如果需要很大，它应该是最快的解决方案;另外，这是最可读的，id说。

Answer 4

这比用于len（测试）= 1000的Amnon变体快两倍：

wanted = (2,4,6)
wanted2 = numpy.expand_dims(wanted, 1)
print test[numpy.any(test[:, 1] == wanted2, 0), :]

从NumPy ndarray中选择行

4 个答案: