pandas/scikit 学习组合列

Question

我有一个包含以下列的 pandas 数据框：

id, cookie_id, file_loc

我正在使用 scikit-learn 进行分层拆分（按id 分层）：

sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=0)
X=df['file_loc']
y=df['id']
for train_index, test_index in sss.split(X, y):
    # print("TRAIN:", train_index, "TEST:", test_index)
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]

但是，我希望将X 视为[df['cookie_id']、df['file_loc']]，以便X_train、y_train 包含信息（cookie_id 和file_loc），但出于某种原因我无法弄清楚如何做到这一点:(

任何指针都会很棒。

Answer 1

我猜你可能正在寻找：

X=df[['cookie_id','file_loc']]