從列表中獲取子列表

Question

我正在嘗試編寫交叉驗證代碼，並且需要遍歷一組並拆分訓練和測試數據。 我試圖以以下方式拆分它，但代碼無法編譯

import numpy as np
X = np.array([[1,2], [3,4], [5,6]])
n_folds = 3

for n in range(n_folds):
    test_fold = X[n]
    train_folds = X[x for x in range(n_folds) if x != n]
    print train_folds, test_fold

預期輸出

[3,4,5,6], [1,2]
[1,2,5,6], [3,4]
[1,2,3,4], [5,6]

有沒有辦法實現這個？

Answer 1

添加額外的括號：

X[[x for x in range(n_folds) if x != n]]

內括號創建一個列表（理解），外括號告訴 numpy 使用該列表進行高級索引。