python - 使用 numpy 过滤元组列表中的行

Question

我正在寻找一种更快的方法来过滤元组列表，使用 numpy 并避免循环。

A = [(27157, 4),
       (24814, 0),
       (1047, 2),
       (18265, 2),
       (2857, 4),
       (23854, 2),
       (36881, 0)]

现在我必须根据第二个元素(即 4)对其进行过滤。如果“4”存在，则形成列表“B”，否则形成列表“C”。

即:

B = [(27157, 4),(2857, 4)]
C = [(24814, 0),(1047, 2),(18265, 2),(23854, 2),(36881, 0)]

Answer 1

通过 NumPy，您可以使用 bool 索引返回数组:

mask = A[:, 1] == 4
B = A[mask]
C = A[~mask]

这要求您的输入是一个 NumPy 数组:

A = np.array([(27157, 4),
              (24814, 0),
              (1047, 2),
              (18265, 2),
              (2857, 4),
              (23854, 2),
              (36881, 0)])