Search code examples
pandassframe

Stacking an array of lists in pandas / SFrame stack equivalent in Pandas


Is there a SFrame stack equivalent in pandas dataframes? Pandas' own stack works only with levels whereas I am looking for expanding a single column at the same level as others which contains lists.

Input Dataframe: There are some more columns like user in actual dataframe

+-------+------------------+
| user  |     friends      |
+-------+------------------+
|  1    |     [2, 3, 4]    |
|  2    |      [5, 6]      |
|  3    | [4, 5, 10, None] |
+----- -+------------------+

Output Dataframe:There are some more columns like user in actual dataframe which should get repeated similarly

+------+--------+
| user | friend |
+------+--------+
|  1   |  2     |
|  1   |  3     |
|  1   |  4     |
|  2   |  5     |
|  2   |  6     |
|  3   |  4     |
|  3   |  5     |
|  3   |  10    |
|  3   |  None  |
+------+--------+

Solution

  • pd.DataFrame.from_items([
        ('user', df.user.values.repeat(df.friends.str.len())),
        ('friends', np.concatenate(df.friends))
    ])
    
       user friends
    0     1       2
    1     1       3
    2     1       4
    3     2       5
    4     2       6
    5     3       4
    6     3       5
    7     3      10
    8     3    None