python-基于列表中的某些元素从数组中选择行

我有一个带值的2d数组，第一个元素表示该行后面数据的ID。例如，对于数组中的一行：[1,45,32,54]，1是键，其余数字是与该键相关的值。在我的数组中，我有编号为1-9的不同键，但我想用ID为的列表中的筛选值创建第二个数组

这是我使用的函数：

def filter_set(dataset, l):
    row=np.where(np.all(dataset[:,0] in l))

    filtered = np.take(dataset,row[0],axis=0)

    return filtered

我这样称呼它：

two_train_set = filter_set(train_set, [0, 1])

我得到了这个错误：

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

---

如果有人能帮助我，我将非常感谢，谢谢。

由于您的密钥范围仅在1到9之间，我建议使用查找表：

>>> def filter_set(dataset, l):
...     mask = np.zeros((10,), dtype=bool)
...     mask[l] = True
...     return dataset[mask[dataset[:, 0]], :]
... 

演示：

您最初的方法也可以通过用np.in1d替换in来工作。请注意，我还用高级索引取代了np.take，我认为这是做同样事情的更现代的方式：

>> def filter_set_OP(dataset, l):
...     row = np.in1d(dataset[:, 0], l)
...     return dataset[row, :]
... 
>>> filter_set_OP(dataset, [5,6,1])
array([[ 6, 44, 93, 41],
       [ 6, 29, 81, 25],
       [ 6, 63, 65, 71],
       [ 1, 69, 45,  2]])

你能详细说明我想用ID为的列表中的筛选值创建第二个数组吗？如果数组较大，还可以添加输入数据集和预期输出，并添加一个较小的版本。我认为in不支持numpy。您可能想尝试np.INAD，并可能摆脱np.all。

>> def filter_set_OP(dataset, l):
...     row = np.in1d(dataset[:, 0], l)
...     return dataset[row, :]
... 
>>> filter_set_OP(dataset, [5,6,1])
array([[ 6, 44, 93, 41],
       [ 6, 29, 81, 25],
       [ 6, 63, 65, 71],
       [ 1, 69, 45,  2]])