返回所选测试集Python的索引

Question

我正在尝试获取测试数据所选择的数据的索引。首先，我对我的数据使用train-test-split

A = [[1,2],[3,4],[6,2],[3,4]]
y = [1,0,0,1]

from sklearn.model_selection import train_test_split
A_train, A_test,y_train,y_test = train_test_split(A, y ,test_size=0.3,random_state=1)

A_test
>> [[3, 4], [6, 2]]

但是，如何获取A中A_test的索引？例如，在这种情况下，它应该返回3，4。非常感谢！

Answer 1

只需在A和y之上使用range(len(A))

from sklearn.model_selection import train_test_split
A_train, A_test, y_train, y_test, index_train, index_test = \
    train_test_split(A, y, range(len(A)))