python - 当列名是整数时合并列

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

所以在这里你可以看到列是 'A' 'B' '1' '2' '5'我想创建一个函数，它接受一个参数作为 int，然后创建一个新的数据帧，其中包含“A”“B”列和作为参数传递的列。所以它会像这样，但是我不确定当其名称是 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

<小时/>

如果列是 str，使用 concat 的解决方案仅转换为 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