CODEWITHSUNDEEP
pythonpandasloops
Jack
Jack

Reputation: 75

Using a loop to select multiple columns from a pandas dataframe

I am trying to create a DataFrame that only has certain columns from a previously created dataframe using a loop.

I have the following dataframe:

      Time  Amount  Amount i=2  Amount i=3  Amount i=4
0    20      10          20          20          20
1    10       5          10          10          10
2    15      25          50          50          50

I am then trying to create a new data frame which has the following columns: Time, Amount, Amount =2, Amount i=3 using a loop function. I understand that this can be solved relatively easily by just select each column, but this is part of a larger project that I can't do that for.

So far I have this:

for i in range (2,4):
    df1 = df[['Time','Amount','Amount i={}'.format(i)]]

But this only pulls out the 'Time' , 'Amount' & 'Amount i=3.

   Time  Amount  Amount i=3
0    20      10          20
1    10       5          10
2     5      25          50

Upvotes: 0

Views: 1459

Answers (2)

mozway
mozway

Reputation: 260500

You could advantageously use filter with a regex:

df1 = df.filter(regex='Time|Amount( i=[23])?$')

'Time|Amount( i=[23])?$' -> Time, or Amount (alone or optionally followed by =x where x is 2 or 3)

output:

   Time  Amount  Amount i=2  Amount i=3
0    20      10          20          20
1    10       5          10          10
2    15      25          50          50

Upvotes: 2

tanmay2893
tanmay2893

Reputation: 123

You can do something like this if you want to use a loop

df1 = df[['Time','Amount'] + ['Amount i={}'.format(i) for i in range(2,4)]]

Upvotes: 2

Related Questions