创建一个python函数来删除数据框列

Question

我的代码如下：

def col_add_drop(dataset, column_name, add = True, **kwargs):
    if add == True:
        column_values = kwargs.get('column_values', [])
        dataset[column_name] = column_values

    else:
        dataset = dataset.drop(column_name, axis=1)

当我尝试添加列时

col_add_drop(dataset = df, column_name = 'income_cat_2', add = True, column_values = df['income_cat'])

这很好用，但是当我尝试通过这样做删除列时：

col_add_drop(dataset = df, column_name = 'income_cat_2', add = False)

它什么都不做，也不显示任何错误。

Answer 1

看起来当您在函数中通过添加列修改它时，pandas 会将其添加到原始列，但是如果您删除列，它只会修改副本。你可以这样做来解决这个问题：

def col_add_drop(dataset, column_name, add = True, **kwargs):
    if add:
        column_values = kwargs.get('column_values', [])
        dataset[column_name] = column_values

    else:
        return dataset.drop(column_name, axis=1)

这样，你必须像这样调用函数：

df = col_add_drop(dataset = df, column_name = 'income_cat_2', add = False)

顺便说一下，if add == True: 相当于 if add:，看起来更干净。