从NumPy ndarray中选择行

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

以下解决方案应该比Amnon的解决方案更快wanted:

# 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只能搜索数组一次(而不是len(wanted)像Amnon的答案那样多次).它还使用Python的内置快速元素查找集,这比列表快得多.它也很快,因为它使用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说.