Complex pandas rows merging keeping order

user_11779593 0
Viewed 29

Let's assume the following pandas dataframe:

df = pd.DataFrame(
    {
        "IDs": [
            ["A","C",],
            ["B"],
            ["A", "B"],
            ["A","C", "D"],
            ["B", "E"],
        ],
        "values": [[1, -3], [1], [2, -2], [4, 1, 3], [2, 5]],
        "timestamp":[
        pd.to_datetime("2022-01-01"),
        pd.to_datetime("2022-01-01"),
        pd.to_datetime("2022-01-02"),
        pd.to_datetime("2022-01-03"),
        pd.to_datetime("2022-01-03"),
    ],
    },
)

    IDs         values      timestamp
0   [A, C]      [1, -3]     2022-01-01
1   [B]         [1]         2022-01-01
2   [A, B]      [2, -2]     2022-01-02
3   [A, C, D]   [4, 1, 3]   2022-01-03
4   [B, E]      [2, 5]      2022-01-03

The idea is to merge rows with a similar timestamp while keeping the IDs sorted and inserting the corresponding values at the right place such that the desired result would be:

    IDs             values          timestamp
0   [A, B, C]       [1, 1, -3]      2022-01-01
1   [A, B]          [2, -2]         2022-01-02
2   [A, B, C, D, E] [4, 2, 1, 3, 5] 2022-01-03

Any idea on how to do that effectively?

FYI: In my application only a very small percentage of rows have duplicated timestamp.

pandas sorting python
1 Answers

You can use .apply() here. Just define a function that takes a sub-dataframe: for each unique timestamp, explodes the lists into rows, sorts by the IDs, and aggregates everything back to a list:

def f(df):
    return (df[['IDs', 'values']]
            .explode(['IDs', 'values'])
            .sort_values('IDs')
            .agg(lambda x: [x.tolist()]))
    
df.groupby('timestamp').apply(f).reset_index(drop=True)
user_146275050
Related