使用可選參數過濾 Numpy 數組

Question

我正在構建一個 function 應該根據輸入准備我的數據。 變量x_imp包含對哪些特征很重要的索引。 但是有時我仍然需要所有功能，所以如果 'x_imp = None' 什么都不會發生。

我的解決方案是這樣的（這不是整個 function 只是輸入）：

def get_train_data(x_cat, x_num,x_imp = None):
        x_cat = x_cat[:,x_imp]
        x_num = x_num[:,x_imp]
    return x_train

但這會改變數據的形狀。 例如，如果data.shape = (4, 5)那么data[:,None].shape = (4, 1, 5)

我該如何避免這個問題？

Answer 1

發生這種情況是因為None切片是np.newaxis的別名。 是否有理由不只添加顯式if語句？

def get_train_data(x_cat, x_num,x_imp = None):
    if x_imp is not None:
        x_cat = x_cat[:,x_imp]
        x_num = x_num[:,x_imp]
    return x_train