根据索引动态创建列表列表-Python

Question

I am trying to make a giant result list with lists of list based on index. I can't predefined how many lists will be inside of the giant list.

id   value
1     30
1     21
1     12
1     0
2     1
2     9
2     14
3     12
3     2
4     3
5     1

result = []
for id, dfs in df.groupby('id'):
    ....

    for i, row in dfs.iterrows():
        x = helper(row[value])
        # If the list is found, append the element
        if (result[i]):
            result[i].append(x)
        # Dynamically make lists base on index
        else:
            result[i] = []

If the list already defined, then just append the value x in the list.

Expected Output:

    first index      second index  third index   fourth index
[[x1,x5,x10,x11,x14], [x2,x4,x9], [x3,x7],       [x20]]

x values are compute by the helper function

Answer 1

It's unclear to me if you want the result as a dataframe or dict with the 'index' as keys or just as a list with the items in the right order. Btw, Python lists start with index 0 .

In [706]: result = collections.defaultdict(list)
     ...: for id, dfs in df.groupby('id'):
     ...:     result[id].extend(list(dfs['value'].values))
     ...:

In [707]: result  # this will be a dict
Out[707]:
defaultdict(list,
            {1: [30, 21, 12, 0], 2: [1, 9, 14], 3: [12, 2], 4: [3], 5: [1]})

In [708]: [result[k] for k in sorted(result.keys())]  # turn it into a list
Out[708]: [[30, 21, 12, 0], [1, 9, 14], [12, 2], [3], [1]]

If you want to apply some operation to each item in the group, like you're doing with helper() , you can do:

In [714]: def helper(val):
     ...:     return 'x' + str(val)  # simplifying whatever helper does

In [715]: result = collections.defaultdict(list)
     ...: for id, dfs in df.groupby('id'):
     ...:     result[id].extend(map(helper, dfs['value'].values))  # pass each value to helper

In [716]: result
Out[716]:
defaultdict(list,
            {1: ['x30', 'x21', 'x12', 'x0'],
             2: ['x1', 'x9', 'x14'],
             3: ['x12', 'x2'],
             4: ['x3'],
             5: ['x1']})

In [717]: [result[k] for k in sorted(result.keys())]
Out[717]:
[['x30', 'x21', 'x12', 'x0'],
 ['x1', 'x9', 'x14'],
 ['x12', 'x2'],
 ['x3'],
 ['x1']]

Note that result[id].extend(...) is not actually needed since each group of values for that 'id' will be passed in together. So you don't need to check if that id already exists in result. It could have been just:

In [720]: result = collections.defaultdict(list)
     ...: for id, dfs in df.groupby('id'):
     ...:     result[id] = list(map(helper, dfs['value'].values))

---

Ideally, you'd want to create helper so that it can be used with pd.apply() , by operating on all the dfs rows together.

Or even better, build helper so that it can do something with the the dataframe of each groupby result, via pd.groupby.GroupBy.apply() .