CODEWITHSUNDEEP
pythonpython-3.xlist
AvePazuzu
AvePazuzu

Reputation: 115

Python: Create sublist based on values/conditions

As a result of a linear program I get a nested list out of a dict which looks like this:

lst = [['a0', 1, 'x_01', 'p_01', 39],
       ['a0', 1, 'x_01', 'p_02', 19],
       ['a0', 1, 'x_02', 'p_01', 10],
       ['a0', 1, 'x_02', 'p_02', 5],
       ['a1', 1, 'x_01', 'y_01', 'p_01', 5],
       ['a1', 1, 'x_01', 'y_01', 'p_02', 5],
       ['a1', 1, 'x_02', 'y_01', 'p_02', 5],
       ['a1', 1, 'x_02', 'y_01', 'p_02', 5]]

Any kind of 'a' indicates a variable which was solved for a range of sets (eg. 1, 'x_01', 'p_01' - etc.) with a value represented by the last element of the list.

Now the task is to create sub lists only containing lists of the same variable. (In the next step they will be passed to a pd.df) The list of lists could look like this or something similar:

new_lst = [[['a0', 1, 'x_01', 'p_01', 39],
       ['a0', 1, 'x_01', 'p_02', 19],
       ['a0', 1, 'x_02', 'p_01', 10],
       ['a0', 1, 'x_02', 'p_02', 5]],
       [['a1', 1, 'x_01', 'y_01', 'p_01', 5],
       ['a1', 1, 'x_01', 'y_01', 'p_02', 5],
       ['a1', 1, 'x_02', 'y_01', 'p_02', 5],
       ['a1', 1, 'x_02', 'y_01', 'p_02', 5]]]

What I've tried so far is: 

st0 = list(set([x[0] for x in lst]))

to get all the unique variables.

And this, but it obviously does not do the job:

n = [x for x in lst for i in st0 if x[0]==i]

What I try here is based on something like this:

n1 = [x for x in lst if x[0]==st0[0]]
n2 = [x for x in lst if x[0]==st0[1]]
ii = [n1,n2]

Based on the model size I have 10, 20 or more different variables of different dimensions. Also the models need to be run several times to test different parameters.

Hard coding is possible, but not desired. If anyone has a good idea please feel free to help me out on this one.

Upvotes: 2

Views: 860

Answers (1)

Austin
Austin

Reputation: 26039

Use itertools.groupby:

from itertools import groupby

lst = [['a0', 1, 'x_01', 'p_01', 39],
       ['a0', 1, 'x_01', 'p_02', 19],
       ['a0', 1, 'x_02', 'p_01', 10],
       ['a0', 1, 'x_02', 'p_02', 5],
       ['a1', 1, 'x_01', 'y_01', 'p_01', 5],
       ['a1', 1, 'x_01', 'y_01', 'p_02', 5],
       ['a1', 1, 'x_02', 'y_01', 'p_02', 5],
       ['a1', 1, 'x_02', 'y_01', 'p_02', 5]]

new_lst = []
s_lst = sorted(lst, key=lambda x: x[0])
for k, v in groupby(s_lst, key=lambda x: x[0]):
    new_lst.append(list(v))

print(new_lst)

Output:

[[['a0', 1, 'x_01', 'p_01', 39], 
  ['a0', 1, 'x_01', 'p_02', 19], 
  ['a0', 1, 'x_02', 'p_01', 10], 
  ['a0', 1, 'x_02', 'p_02', 5]], 
 [['a1', 1, 'x_01', 'y_01', 'p_01', 5], 
  ['a1', 1, 'x_01', 'y_01', 'p_02', 5], 
  ['a1', 1, 'x_02', 'y_01', 'p_02', 5], 
  ['a1', 1, 'x_02', 'y_01', 'p_02', 5]]]

Upvotes: 4

Related Questions