Pandas add columns to a dataframe

Question

So, I've seen this answer here, which is sensible for functions which return one output. What if my function has multiple outputs?

More concretely, let's say I am cross-referencing some data on some ID. But when I call certain IDs, it returns multiple matches, which I want to put into different columns.

An example of this would be something like the below, where worker 3 has two bosses, 0 and 2, while worker 1 has one boss, 2.

Worker_ID Boss_ID
        3       0
        3       2
        1       2

Is it possible to create the second column and populate without first going through, counting the number of matches and creating the relevant number of columns?

EDIT:

I'd like something like this in short-form:

Worker_ID  Boss_ID_1 Boss_ID_2   ...as necessary
        3          0         2
        1          2       nan

Looks like you need `merge`. Can you add more details of the data you are dealing with? An example of what you have and what you want would be great — pomber, May 03 '18 at 21:48
can workers only have either 1 or 2 bosses? or could they have more? — Simon, May 03 '18 at 22:03

score 0 · Answer 1 · answered May 03 '18 at 22:03

Create a key by using cumcount then we can using pivot

df.assign(key=df.groupby('Worker_ID').cumcount()+1).\
   pivot(index='Worker_ID',columns='key',values='Boss_ID').\
      add_prefix('Boss_ID_')
Out[242]: 
key        Boss_ID_1  Boss_ID_2
Worker_ID                      
1                2.0        NaN
3                0.0        2.0

Pandas add columns to a dataframe

1 Answers1