从 NumPy ndarray 中选择行

我只想从 a 中选择某些行NumPy http://en.wikipedia.org/wiki/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 本身是否有更好的方法来做到这一点？

以下解决方案应该比 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仅数组once（而不是尽可能多len(wanted)次如阿姆农的回答）。它还使用Python内置的快速元素查找sets，这比列表快得多。它也很快，因为它使用 Numpy 的快速循环。您还可以获得以下优化in操作员：一次wanted如果元素匹配，则不必测试其余元素（与 Amnon 的“逻辑或”方法相反，如果所有元素都在wanted无论如何都经过测试）。

或者，您可以使用以下单行代码，它也只遍历您的数组一次：

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

不过，这要慢得多，因为这会在每次迭代时提取第二列中的元素（而不是像该答案的第一个解决方案那样一次性完成）。