python - 选择 numpy 数组，使最后一个元素为 1

Question

假设我们有一个像这样的数据集:

X = 
6  2   1
-2 4  -1
4  1  -1
1  6   1
2  4   1
6  2   1

我想从中获取两个数据，一个的最后一位数字为 1，另一个的最后一位数字为 -1。

X0 = 
-2 4  -1
4  1  -1

并且，

X1 = 
6  2   1
1  6   1
2  4   1
6  2   1

我们如何在 numpy 中有效地做到这一点？

在简单的Python中，我可以这样做:

dataset = np.loadtxt('data.txt')

X0, X1 = [], []
for i in range(len(X)):
    if X[i][-1] == 1:
        X0.append(X[i])
    else:
        X1.append(X[i])

这又慢又麻烦，Numpy 又快又简单，所以，如果 numpy 有更简单的方法，我将不胜感激。谢谢。

Answer 1

假设你有一个数组:

>>> arr
array([[ 6,  2,  1],
       [-2,  4, -1],
       [ 4,  1, -1],
       [ 1,  6,  1],
       [ 2,  4,  1],
       [ 6,  2,  1]])

然后简单地:

>>> mask1 = arr[:, -1] == 1
>>> mask2 = arr[:, -1] == -1
>>> X1 = arr[mask1]
>>> X2 = arr[mask2]

结果:

>>> X1
array([[6, 2, 1],
       [1, 6, 1],
       [2, 4, 1],
       [6, 2, 1]])
>>> X2
array([[-2,  4, -1],
       [ 4,  1, -1]])