2

Let us say I have the following column names in pandas: ["A", "B"]

My problem is that I want to use a for loop that grabs the column names from the list and creates a new column name that includes part of those elements from the list.

In each iteration I would like to create the following:

a = ["A", "B"]
for elements i in a:
    elements + "_c" = SOME FUNCTION column <df.elements>

That should result in the addition of two new columns with names A_c and B_c.

vestland
  • 55,229
  • 37
  • 187
  • 305
Guga
  • 349
  • 1
  • 4
  • 14

1 Answers1

2

You can use Series.apply in loop:

a = ["A", "B"]
for i in a:
    df[i + "_c"] =  df[i].apply(SOME FUNCTION)

Or DataFrame.apply with add_suffix for new df and then join to original:

df1 = df[a].apply(SOME FUNCTION).add_suffix('_c')
df = df.join(df1)

Sample:

df = pd.DataFrame({'A':[4,5,4,5,5,4],
                   'B':[7,8,9,4,2,3],
                   'C':[1,3,5,7,1,0]})

print (df)
   A  B  C
0  4  7  1
1  5  8  3
2  4  9  5
3  5  4  7
4  5  2  1
5  4  3  0

def FUNCTION(x):
    return x + 10

a = ["A", "B"]

for i in a:
    df[i + "_c"] =  df[i].apply(FUNCTION) 

print (df)
   A  B  C  A_c  B_c
0  4  7  1   14   17
1  5  8  3   15   18
2  4  9  5   14   19
3  5  4  7   15   14
4  5  2  1   15   12
5  4  3  0   14   13

df = df.join(df[a].apply(FUNCTION).add_suffix('_c'))
print (df)
   A  B  C  A_c  B_c
0  4  7  1   14   17
1  5  8  3   15   18
2  4  9  5   14   19
3  5  4  7   15   14
4  5  2  1   15   12
5  4  3  0   14   13
jezrael
  • 822,522
  • 95
  • 1,334
  • 1,252
  • thank you! It works fine. My problem was that I had to specify a parameter inside the FUNCTION that would has the name containing the names of columns. – Guga Jul 22 '17 at 03:54
  • It depends what need - if need working in series in function use `x` and if need column name use `x.name` in function `FUNC` – jezrael Jul 22 '17 at 04:01