Merging elements in 2D list based on common first elements

Question

Given the following list:

lst = [[3,5],[3,10],[3,15],[3,102],[5,21],[5,23],[5,50]]

I want to obtain the following [[3,5,10,15,102], [5,21,23,50]]

Note that the list is sorted in ascending order based on the value of the first element.

What would be the most efficient way of doing this? This is what I was thinking:

Step 1: Create a list with unique first elements. (i.e. 3 and 5 in this case)

first_elements = [] #initialize empty list to which we will append all first elements
for i in range(len(lst)):
    first_elements.append(lst[i][0])
first_elements = list(set(first_elements)) #Filter out the unique first elements

    first_elements = [3,5]

Step 2: Filter lst based on the first element. Append these to a new list.

new_merged_list = [] # create new list to append to
for i in range(len(first_elements)): 
    first_element_to_filter_by = first_elements[i]
    filtered_2d_list           = [i for i in lst if i[0] == first_element_to_filter_by]
    new_merged_list.append([first_element_to_filter_by])

    for j in range(len(filtered_2d_list)):
        (new_merged_list[i]).append(filtered_2d_list[j][1])    

This gives me the correct answer, as shown below.

new_merged_list = [[3, 5, 10, 15, 102], [5, 21, 23, 50]]

My question - is there a more efficient way to do this? I don't know how well this would scale to a list that is (for instance) 100000 x 2.

Appreciate the help!

Answer 1

You can use defaultdict here. This will work even if your lst is not sorted.

from collections import defaultdict
new = defaultdict(list)
lst = [[3, 5], [3, 10], [3, 15], [3, 102], [5, 21], [5, 23], [5, 50]]

for k,v in lst:
    new[k].append(v)

new = [[k]+v for k,v in new.items()]
# [[3, 5, 10, 15, 102], [5, 21, 23, 50]]
# Or 
new = [[k,*v] for k,v in new.items()]
# [[3, 5, 10, 15, 102], [5, 21, 23, 50]]

Answer 2

You could do something like this (without importing a module)

first_elements = list(set([item[0] for item in lst]))
result = [[elem] for elem in first_elems]
for sublist in lst:
    result[first_elems.index(sublist[0])].append(sublist[1])