Python 列名为整数时合并列

这是我的数据帧：

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”列，并且该列作为arumment传递。 所以它是这样的，但是我不知道当它的名字是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

---

最简单的解决方案是按子集选择：

---

如果列为

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

您还可以嵌入

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

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