用子列表的给定元素将列表中的子列表集分开

Question

Input is:

 main_list = [['a','1','x'], ['b','2', 'm'], ['a', '23', 'm'], ['c', '34', 'x'], ['b', '11'], ['b', '32', 'y']]

Final result should be:

 a_list = [['a','1','x'], ['a', '23', 'm']]
 b_list = [['b','2', 'm'], ['b', '32', 'y']]   <-- note that list ['b', '11'] has not been included as it has less elements.
 c_list = ['c', '34', 'x']

Personally, I did something like this:

 a_list = [[p for p in item] for item in main_list if "a" in item[0]]
 b_list = [[p for p in item] for item in main_list if "b" in item[0] and len(item) > 2]
 c_list = [[p for p in item] for item in main_list if "c" in item[0]]

Wondering if there is a way to extract all three lists at once instead of walking main_list three times.

To extend the problem, how would you do it if no element is given but your algorithm should sort the sublists by common elements at, let's say, a given position.

Cheers,

Answer 1

Sometimes the best way is the simplest way. List comprehensions are good, but can get unwieldy and unreadable. They are also no faster than a comparable for loop. The reason you can't convert this into a list comprehension easily is that you will need to append to existing lists to convert 1 list into 3. You might be able to do something using yield but that again is quite a bit of work.

main_list = [['a','1','x'], ['b','2', 'm'], ['a', '23', 'm'], ['c', '34', 'x'], ['b', '11'], ['b', '32', 'y']]

a_list, b_list, c_list = [],[],[]

for item in main_list:
    if "a" in item[0]:
        a_list.append(item)
    elif "b" in item[0] and len(item) > 2:
        b_list.append(item)
    elif "c" in item[0]:
        c_list.append(item)

print a_list
print b_list
print c_list

Outputs:

[['a', '1', 'x'], ['a', '23', 'm']]
[['b', '2', 'm'], ['b', '32', 'y']]
[['c', '34', 'x']]

Answer 2

groupby in the itertools module will get the information you are looking for (and more):

In [31]: main_list = [['a','1','x'], ['b','2', 'm'], ['a', '23', 'm'], ['c', '34', 'x'], ['b', '11'], ['b', '32', 'y']]

In [32]: g = itertools.groupby(main_list, key=lambda x: (x[0], len(x)))

To see what is in that generator, run:

In [33]: for x in g: print x[0], list(x[1])
('a', 3) [['a', '1', 'x']]
('b', 3) [['b', '2', 'm']]
('a', 3) [['a', '23', 'm']]
('c', 3) [['c', '34', 'x']]
('b', 2) [['b', '11']]
('b', 3) [['b', '32', 'y']]

As you can see, this produces information for all lengths. You are only interested in length=3. Depending on how you actually want to use the generator, there are a variety of ways to throw away the other lengths.

As an example, this extracts the information from the generator g into a dictionary d with the results that you were looking for:

In [65]: g = itertools.groupby(main_list, key=lambda x: (x[0], len(x)))

In [66]: d = collections.defaultdict(list)

In [67]: for x in g:
    if x[0][1]==3: d[x[0][0]].append(list(x[1])[0])
   ....:     

In [68]: d['a']
Out[68]: [['a', '1', 'x'], ['a', '23', 'm']]

In [69]: d['b']
Out[69]: [['b', '2', 'm'], ['b', '32', 'y']]

In [70]: d['c']
Out[70]: [['c', '34', 'x']]