10

I'm giving a toy example but it will help me understand what's going on for something else I'm trying to do. Let's say I want a new column in a dataframe 'optimal_fruit' that is apples * orange - bananas.

I can do something like this to get it. 

df2['optimal_fruit'] = df2['apples'] * df2['oranges'] - df2['bananas'] 


apples  oranges bananas optimal_fruit
1       6       11      -5
2       7       12      2
3       8       13      11
4       9       14      22
5       10      15      35

What is happening if I try to do something like this? And how could I do this in a list comprehension?

df2['optimal_fruit'] = [x * y - z for x in df2['apples'] for y in df2['oranges'] for z in df2['bananas']]

I get an error of:

ValueError: Length of values does not match length of index

As always, thank you all so much for your help!

WhitneyChia
  • 746
  • 2
  • 11
  • 28
  • 1
    possible duplicate of https://stackoverflow.com/questions/58567199/memory-efficient-way-for-list-comprehension-of-pandas-dataframe-using-multiple-c/62064720#62064720 and https://stackoverflow.com/questions/52607864/pandas-list-comprehension-tuple-from-dataframe/62064822#62064822 But this was the first question, so the possible duplicates are rather the 2 links. – questionto42 May 28 '20 at 13:03
  • Does this answer your question? [Selecting multiple columns in a Pandas dataframe](https://stackoverflow.com/questions/11285613/selecting-multiple-columns-in-a-pandas-dataframe) – questionto42 Jul 08 '22 at 18:03

4 Answers4

27

Essentially your list comprehension statement is a set of 3 nested loops. In code:

l = []
for x in df2['apples']:
    for y in df2['oranges']:
        for z in df2['bananas']:
            l.append(x * y - z)

The length of your resultant list will be power-of-3 times the length of your DataFrame (5x5x5 = 125). Hence the error. To fix, you need the equivalent of:

for x, y, z in zip(df2['apples'], df2['oranges'], df2['bananas']):
    l.extend([x * y - z])

In terms of list comprehension:

[x * y - z for x, y, z in zip(df2['apples'], df2['oranges'], df2['bananas'])]
Nas Banov
  • 28,347
  • 6
  • 48
  • 67
Kartik
  • 8,347
  • 39
  • 73
4

The reason why your new method doesn't work is because the list comprehension produces data that is longer than the number of indices in your dataframe. A quick fix for that would be something like:

[x * y - z for x,y,z in zip(df2['apples'], df2['oranges'], df2['bananas'])]
jtitusj
  • 3,046
  • 3
  • 24
  • 40
2

You can get all the values of the row as a list using the np.array() function inside your list of comprehension.

The following code solves your problem:

df2['optimal_fruit'] = [x[0] * x[1] - x[2] for x in np.array(df2)]

It is going to avoid the need of typing each column name in your list of comprehension.

Adrian Mole
  • 49,934
  • 160
  • 51
  • 83
Tiago Pitchon
  • 63
  • 6
1

If you do not want to repeat df2 for each column:

[row[0][0]*row[0][1]-row[0][2] for row in zip(df2[['apples', 'oranges', 'bananas']].to_numpy())]

or

def func(row):
    print(row[0]*row[1]-row[2])

[func(*row) for row in zip(df2[['apples', 'oranges', 'bananas']].to_numpy())]

Further reading:

EDIT:

Please use df.iloc and df.loc instead of df[[...]], see Selecting multiple columns in a Pandas dataframe

questionto42
  • 7,175
  • 4
  • 57
  • 90