Question

这是我的数据框：

A   B     1    2    5

1   1     26    5   5     
2   6     4     2   2    
2   3     8     7   7     
2   4     3    10   19

所以在这里你可以看到列是＆＃39; A＆＃39; ＆＃39; B＆＃39; ＆＃39; 1＆＃39; ＆＃39; 2＆＃39; ＆＃39; 5＆＃39; 我想创建一个接受一个参数作为int的函数，然后创建一个具有＆＃39; A＆＃39;的新数据框。＆＃39; B＆＃39;作为arugment传递的列和列。所以它会像这样，但我不知道如何获取列的名称是一个int任何想法吗？谢谢！

def merge(arg):
    new_df=pd.concat([df.A, df.B, df.arg], axis=1 )
merge(5)

A   B     5

1   1     5   
2   6     2  
2   3     7   
2   4     19

Answer 1

最简单的解决方案是使用copy

按子集进行选择

def merge(arg):
     return df[['A', 'B', str(arg)]].copy()

print (merge(5))
   A  B   5
0  1  1   5
1  2  6   2
2  2  3   7
3  2  4  19

concat的解决方案，如果列str仅转换为str - 效果不错，但速度较慢：

def merge(arg):
     return pd.concat([df.A, df.B, df[str(arg)]], axis=1 )

print (merge(5))
   A  B   5
0  1  1   5
1  2  6   2
2  2  3   7
3  2  4  19

如果int：

def merge(arg):
     return pd.concat([df.A, df.B, df[arg]], axis=1 )

print (merge(5))
   A  B   5
0  1  1   5
1  2  6   2
2  2  3   7
3  2  4  19

Answer 2

def merge(arg):
    return df.loc[:, ['A', 'B', str(arg)]]

merge(2)

   A  B   2
0  1  1   5
1  2  6   2
2  2  3   7
3  2  4  10

您还可以嵌入if/else逻辑以使其更强大

def merge(arg):
    arg = arg if arg in df.columns else str(arg)
    return df.loc[:, ['A', 'B', arg]]

print(merge(2))

   A  B   2
0  1  1   5
1  2  6   2
2  2  3   7
3  2  4  10

或者有点聪明的交集

def merge(arg):
    cols = df.columns.intersection(['A', 'B', arg, str(arg)])
    return df.loc[:, cols]

print(merge(2))

   A  B   2
0  1  1   5
1  2  6   2
2  2  3   7
3  2  4  10

列名是整数时合并列

2 个答案: