Search code examples
pandasdataframeseries

How do I remap a dataframe with a dict containing nested list using map function?


I have a dataframe like this:

    ID
0   123
1   123
2   123
3   123
4   123
5   123
6   456
7   456
8   456
9   123

I want to add a new column with map function using dict d:

d = {'123': [1, 2, 3, 1, 2, 1, 5], '456': [1, 2, 1]}

Expected output:

    ID  Count
0   123   1
1   123   2
2   123   3
3   123   1
4   123   2
5   123   1
6   456   1
7   456   2
8   456   1
9   123   5

But df.ID.map(d) returns:

0    [1, 2, 3, 1, 2, 1, 5]
1    [1, 2, 3, 1, 2, 1, 5]
2    [1, 2, 3, 1, 2, 1, 5]
3    [1, 2, 3, 1, 2, 1, 5]
4    [1, 2, 3, 1, 2, 1, 5]
5    [1, 2, 3, 1, 2, 1, 5]
6             [1, 2, 1]
7             [1, 2, 1]
8             [1, 2, 1]
9    [1, 2, 3, 1, 2, 1, 5]

Solution

  • You can use groupby+apply:

    df.groupby('ID').apply(lambda g: pd.Series(d[g.name]))
    

    Example:

    >>> df['Count'] = df.groupby('ID').apply(lambda g: pd.Series(d[g.name])).to_list()
    >>> df
        ID  Count
    0  123      1
    1  123      2
    2  123      3
    3  123      1
    4  123      2
    5  123      1
    6  456      1
    7  456      2
    8  456      1
    

    edit. variant for unordered input:

    (df.join(df.groupby('ID').apply(lambda g: pd.Series(d[g.name],
                                                        name='Count',
                                                        index=g.index))
               .droplevel(0))
    )
    

    output:

        ID  Count
    0  123      1
    1  123      2
    2  123      3
    3  123      1
    4  123      2
    5  123      1
    6  456      1
    7  456      2
    8  456      1
    9  123      5