提取按条件过滤的测试集的子集

Question

我已将数据集分为训练和测试：

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0)

然后我计算了一个预测变量和相关分数。我现在想看看这个预测器如何处理满足特定条件的数据子集。

我需要提取X_test的子集和相应的y_test。

用于说明目的

X_test = [[9.3400e+01 9.4710e+01 1.2100e+03 0.0000e+00 0.0000e+00 0.0000e+00 1.0000e+00]
          [9.1210e+01 9.1890e+01 1.3600e+03 0.0000e+00 1.0000e+00 0.0000e+00 0.0000e+00]
          [8.8810e+01 9.0790e+01 1.3400e+03 0.0000e+00 0.0000e+00 0.0000e+00 1.0000e+00]]
y_test = [0 1 0]

及条件：

 (X_test[i][4]== 1)

我应该最终得到

X_test_cond = [[9.1210e+01 9.1890e+01 1.3600e+03 0.0000e+00 1.0000e+00 0.0000e+00 0.0000e+00]]
y_test_cond = [1]

如何过滤 X_test 并保留相应的索引以创建关联的 y_test 子集？

Answer 1

所以这是一个答案...但这纯粹是程序性的...

X_test_tmp = []
y_test_tmp = []
for i in range(len(X_test)):
    if X_test[i][4] == 1:
        X_test_tmp.append(X_test[i])
        y_test_tmp.append(y_test[i])
X_test_cond = np.array(X_test_tmp)
y_test_cond = np.array(y_test_tmp)

我们可以用更少的行数来完成吗？

Answer 2

X_test_cond = X_test[y_test == 1] 

返回 X_test 中与 y_test 等于 1 的索引相对应的实例。