Merging sublists into a list based on sublist item in python

Question

Input:

I have this ordered list.

[[1, 'A'], [1, 'B'],[1, 'D'], [2, 'A'],[2,'D'], [3, 'C'], [4, 'D'], [5, 'B'], [6, 'D']]

Desired output

[[1,['A','B','D']],[2, ['A','D']], [3, 'C'], [4, 'D'], [5, 'B'], [6, 'D']]

Since the first item of those two sublists are same.

Also can i convert into a dictionary with a key and those values pair. Like

{1:['A','B','D'],2:['A','D'],3:['C']}

What is the easiest and simplest way to do this?

Answer 1

If the data is ordered, then itertools.groupby is a good approach:

>>> from itertools import groupby
>>> from operator import itemgetter
>>> data = [[1, 'A'], [1, 'B'], [2, 'A'], [3, 'C'], [4, 'D'], [5, 'B'], [6, 'D']]
>>> final_data = []
>>> final_data = []
>>> for k, g in groupby(data, itemgetter(0)):
...     group = list(g)
...     if len(group) == 1:
...         final_data.append(group[0])
...     else:
...         final_data.append([k, [sub[1] for sub in group]])
...
>>> final_data
[[1, ['A', 'B']], [2, 'A'], [3, 'C'], [4, 'D'], [5, 'B'], [6, 'D']]
>>>

If you want the results in a dictionary, that is even easier:

>>> grouped_dict = {}
>>> for num, letter in data:
...     grouped_dict.setdefault(num, []).append(letter)
...
>>> grouped_dict
{1: ['A', 'B'], 2: ['A'], 3: ['C'], 4: ['D'], 5: ['B'], 6: ['D']}
>>>

Answer 2

You can create the dictionary directly from the input.

from collections import defaultdict

input = [[1, 'A'], [1, 'B'],[1, 'D'], [2, 'A'],[2,'D'], [3, 'C'], [4, 'D'], [5, 'B'], [6, 'D']]

d = defaultdict(list)
for el in input: d[el[0]].append(el[1])

The output of d will be:

{1: ['A', 'B', 'D'], 2: ['A', 'D'], 3: ['C'], 4: ['D'], 5: ['B'], 6: ['D']}