Generate conditional lists of lists in Pandas, "Pythonically"

user_9870343 0
Viewed 16

I want to generate a conditional list of lists. The number of embedded lists is determined by the number of unique conditions, and each embedded list contains values from a given condition.

I can generate this list of lists using a for-loop. See the code below. However, I am looking for a faster and more Pythonic (i.e, no for-loop) approach.

import pandas as pd
from random import randint

example_conditions = ["A","A","B","B","B","C","D","D","D","D"]
example_values = [randint(-100,100) for _ in example_conditions ]

df = pd.DataFrame({
    "conditions":example_conditions,
    "values": example_values
})

lol = []
for condition in df["conditions"].unique():
    sublist = df.loc[df["conditions"]==condition]["values"].values.tolist()
    lol.append(sublist)

Thanks!

pandas python-3.x group-by
1 Answers

Try:

x = df.groupby("conditions")["values"].agg(list).to_list()
print(x)

Prints:

[[-1, 78], [33, 74, -79], [59], [-32, -2, 52, -66]]

Input dataframe:

  conditions  values
0          A      -1
1          A      78
2          B      33
3          B      74
4          B     -79
5          C      59
6          D     -32
7          D      -2
8          D      52
9          D     -66
user_100359850
Related