选择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快速而简单所以，如果有更简单的方法，我会很感激。感谢。

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]])

Answer 2

您可以使用numpy并使用切片来访问您的数据，例如：

X[X[:, 2] == 1]  # Returns all rows where the third column equals 1

或作为一个完整的例子：

import numpy as np

# Random data set
X = np.zeros((6, 3))
X[:3, 2] = 1
X[3:, 2] = -1
np.random.shuffle(X)

print(X[X[:, 2] == 1])
print('-')
print(X[X[:, 2] == -1])

Answer 3

import numpy as np
x = np.array(x)
x0 = x[np.where(a[:,2]==-1)]
x1 = x[np.where(a[:,2]==1)]