Python 返回所选测试集的索引

我正在尝试获取由测试数据选择的数据的索引。首先，我对数据使用列车测试拆分

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_测试的索引？例如，在本例中，它应该返回3，4。非常感谢

只需在

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