Распространенный способ выбора столбцов в массиве numpy и пандах

Question

Мне нужно написать объект, который принимает в качестве входных данных кадр данных pandas или массив numpy (аналогично поведению sklearn). В одном из методов для этого объекта мне нужно выбрать столбцы (не конкретный фиксированный, я получаю несколько индексов столбцов на основе других вычислений).

Итак, чтобы сделать мой код совместимым с обоими типами ввода, я попытался найти общий способ выбора столбцов и попробовал такие методы, как X[:,0] (не работает на фреймах данных pandas), 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 dataframe, col_indices - это индексы столбцов, которые вы хотите выбрать.