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

Question

I have to write an object that takes either a pandas data frame or a numpy array as the input (similar to sklearn behavior). In one of the methods for this object, I need to select the columns (not a particular fixed one, I get a few column indices based on other calculations).

So, to make my code compatible with both input types, I tried to find a common way to select columns and tried methods like X[:,0] (doesn't work on pandas dataframes), X[0] and others but they select differently. Is there a way to select columns in a similar fashion across pandas and numpy?

If no then how does sklearn work across these data structures?

Answer 1

You can use an if condition within your method and have separate selection methods for pandas dataframes and numpy arrays. Given sample code below.

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]

Here, var is your input which can be a numpy array or pandas dataframe, col_indices are the indices of the columns you want to select.