python - 基于行比较的 2D ndarray bool 索引

Question

我有一对相同维度的二维数组 (n, 3)。我想根据第二个索引从第一个中进行选择。我的想法如下:

data[labels == row]

其中row是长度为3的向量。内部 bool 比较给出形状为 (n, 3) 的数组。索引给出了一个平面一维数组。

我的问题是，我必须手动 reshape 数组，或者在数组labels == row上使用类似np.all的东西。

如果 data 是 pandas DataFrame，这实际上可以正常工作。使用纯 ndarray 执行此操作的正确方法是什么？

Answer 1

使用(labels == row).all(axis=1)选择所有值匹配的行:

import numpy as np
np.random.seed(2016)

labels = np.random.randint(10, size=(10, 3))
data = np.random.randint(10, size=(10, 3))
# array([[0, 8, 2],
#        [3, 2, 2],
#        [4, 0, 9],
#        [0, 4, 9],
#        [5, 5, 1],
#        [7, 8, 0],
#        [0, 9, 5],
#        [0, 6, 2],
#        [0, 0, 5],
#        [5, 0, 7]])

row = labels[::3] = labels[0]
data[(labels == row).all(axis=1)]

产量

array([[0, 8, 2],
       [0, 4, 9],
       [0, 9, 5],
       [5, 0, 7]])

<小时/>

请注意， bool 数组 labels == row 有一些 True 值在不完全匹配的行上:

In [138]: labels == row
Out[138]: 
array([[ True,  True,  True],
       [ True, False, False],    # <-- a lone True value
       [False,  True, False],    # <--
       [ True,  True,  True],
       [False, False, False],
       [False, False,  True],    # <--
       [ True,  True,  True],
       [False, False, False],
       [False, False, False],
       [ True,  True,  True]], dtype=bool)

因此，data[labels == row] 返回一些与完整行匹配不相关的值:

In [141]: data[labels == row]
Out[141]: array([0, 8, 2, 3, 0, 0, 4, 9, 0, 0, 9, 5, 5, 0, 7])
                          ^  ^           ^
                          |  |           |
                          not related to a complete row match