CODEWITHSUNDEEP
listfunctionpandas
user36729
user36729

Reputation: 575

Getting a list as the result of a function in pandas

I have data frame in pandas and I have written a function to use the information in each row to generate a new column. I want the result to be in a list format:

      A    B    C
      3    4    1
      4    2    5

     def Computation(row):
          if row['B'] >= 3:
              return [s for s in range(row['C'],50)]
          else:
              return [s for s in range(row['C']+2,50)]

     df['D'] = df.apply(Computation, axis = 1)

However, I am getting the following error:

"could not broadcast input array from shape (308) into shape (9)"

Could you please tell me how to solve this problem?

Upvotes: 1

Views: 64

Answers (1)

Ami Tavory
Ami Tavory

Reputation: 76346

Say you start with

In [25]: df = pd.DataFrame({'A': [3, 4], 'B': [4, 2], 'C': [1, 5]})

Then there are at least two ways to do it.

You can apply twice on the C column, but switch on the B column:

In [26]: np.where(df.B >= 3, df.C.apply(lambda c: [s for s in range(c, 50)]), df.C.apply(lambda c: [s for s in range(c + 2, 50)]))
Out[26]: 
array([ [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49],
       [7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49]], dtype=object)

Or you can apply on the entire row and switch on the B value per row:

In [27]: df.apply(lambda r: [s for s in range(r.C, 50)] if r.B >= 3 else [s for s in range(r.C + 2, 50)], axis=1)
Out[27]: 
0    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14...
1    [7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, ...

Note that the return types are different, but, in each case, you can still write

df['foo'] = <each one of the above options>

Upvotes: 1

Related Questions