在numpy数组和pandas数据框中选择列的常用方法

Question

我必须编写一个以pandas数据框或numpy数组作为输入的对象（类似于sklearn行为）。在用于该对象的一种方法中，我需要选择列（不是特定的固定列，我会根据其他计算获得一些列索引）。

因此，为了使我的代码与两种输入类型兼容，我尝试找到一种选择列的常用方法，并尝试了X[:,0]（不适用于熊猫数据框），X[0]和其他，但他们选择不同。有没有办法在大熊猫和numpy中以类似的方式选择列？

如果否，那么sklearn在这些数据结构中如何工作？

Answer 1

您可以在方法中使用if条件，并对pandas数据框和numpy数组具有单独的选择方法。在下面给出示例代码。

def method_1(self, var, col_indices):
    if isinstance(var, pd.DataFrame):
        selected_columns = var[var.columns[col_indices]]
    else:
        selected_columns = var[:,col_indices]

在这里，var是您的输入，可以是numpy数组或pandas数据框，col_indices是您要选择的列的索引。