Divide list to multiple lists based on elements value

Question

I have the following list:

initial_list = [['B', 'D', 'A', 'C', 'E']]

On each element of the list I apply a function and put the results in a dictionary:

for state in initial_list:
    next_dict[state] = move([state], alphabet)

This gives the following result:

next_dict = {'D': ['E'], 'B': ['D'], 'A': ['C'], 'C': ['C'], 'E': ['D']}

What I would like to do is separate the keys from initial_list based on their values in the next_dict dictionary, basically group the elements of the first list to elements with the same value in the next_dict:

new_list = [['A', 'C'], ['B', 'E'], ['D']]

'A' and 'C' will stay in the same group because they have the same value 'C', 'B' and 'D' will also share the same group because their value is 'D' and then 'D' will be in it's own group.

How can I achieve this result?

Answer 1

You need groupby, after having sorted your list by next_dict values :

It generates a break or new group every time the value of the key function changes (which is why it is usually necessary to have sorted the data using the same key function).

from itertools import groupby

initial_list = ['B', 'D', 'A', 'C', 'E']

def move(letter):
    return {'A': 'C', 'C': 'C', 'D': 'E', 'E': 'D', 'B': 'D'}.get(letter)
sorted_list = sorted(initial_list, key=move)
print [list(v) for k,v in groupby(sorted_list, key=move)]
#=> [['A', 'C'], ['B', 'E'], ['D']]