Sort list into list based on common string while keeping entire list

user_16503179 0
Viewed 79

I have a dataset that sorta looks like the following:

old_data = [[0,0,0,X],[1,0,0,A],[2,0,0,HJ],[3,0,0,X],[4,0,0,A],[5,0,0,K],[6,0,0,K]

I want to be able to sort the lists into their own seperate lists which would like the following:

X_list = [[0,0,0,X],[3,0,0,X]]
A_list = [[1,0,0,A],[4,0,0,A]]
HJ_list = [[2,0,0,HJ]]
K_list = [[5,0,0,K],[6,0,0,K]]

I tried playing around with a few features in pandas and some other solutions on stack, however, they all group the specific value into one list, which ends up deleting all the other data in the list. Furthermore, there will be an X amount of list[3], therefore it will need to generate X amount of lists. The previous thing I tried was this:

values = sorted(set(map(lambda x: x[3], old_data)))
newlist = [[y[3] for y in old_data if y[3] == x] for x in values]

which outputs, something as follows:

newlist = [[X,X],[A,A],[HG],[K,K]]

Possibly if there was someway to rearrange the code above maybe to keep all the data and create an X amount of lists with a for loop? I don't really understand the code above, so I couldn't figure out how to do it. Otherwise, possibly some way using pandas?

pandas-groupby pandas dataframe list python
3 Answers

You can simply use itertools.groupby:

from itertools import groupby
old_data = [[0,0,0,'X'],
            [1,0,0,'A'],
            [2,0,0,'HJ'],
            [3,0,0,'X'],
            [4,0,0,'A'],
            [5,0,0,'K'],
            [6,0,0,'K']]

{k: list(v) for k,v in groupby(sorted(old_data, key=lambda x: x[-1]), lambda x: x[-1])}

output:

{'A': [[1, 0, 0, 'A'], [4, 0, 0, 'A']],
 'HJ': [[2, 0, 0, 'HJ']],
 'K': [[5, 0, 0, 'K'], [6, 0, 0, 'K']],
 'X': [[0, 0, 0, 'X'], [3, 0, 0, 'X']]}

and if you want a custom group name:

{k+'_list': list(v)
 for k,v in groupby(sorted(old_data, key=lambda x: x[-1]),
lambda x: x[-1])}

{'A_list': [[1, 0, 0, 'A'], [4, 0, 0, 'A']],
 'HJ_list': [[2, 0, 0, 'HJ']],
 'K_list': [[5, 0, 0, 'K'], [6, 0, 0, 'K']],
 'X_list': [[0, 0, 0, 'X'], [3, 0, 0, 'X']]}
user_163434640

Assuming your A,X etc are strings, we can make a dict, with the keys as your labels, and the values as your lists.

We loop over your data, and use the .setdefault method to mkae an empty list for any keys we havent seen yet. Then we append your list on

old_data = [[0,0,0,'X'],[1,0,0,'A'],[2,0,0,'HJ'],[3,0,0,'X'],[4,0,0,'A'],[5,0,0,'K'],[6,0,0,'K']]

output = {}
for i in old_data:
    output.setdefault(i[3]+'_list', []).append(i)
output

{'X_list': [[0, 0, 0, 'X'], [3, 0, 0, 'X']],
 'A_list': [[1, 0, 0, 'A'], [4, 0, 0, 'A']],
 'HJ_list': [[2, 0, 0, 'HJ']],
 'K_list': [[5, 0, 0, 'K'], [6, 0, 0, 'K']]}
user_37609200

Create a Series from your list of lists and then use groupby to split them based on the last element.

import pandas as pd

s = pd.Series(old_data)
s.groupby(s.str[-1]).agg(list)

A     [[1, 0, 0, A], [4, 0, 0, A]]
HJ                 [[2, 0, 0, HJ]]
K     [[5, 0, 0, K], [6, 0, 0, K]]
X     [[0, 0, 0, X], [3, 0, 0, X]]
dtype: object
user_43333590
Related