Home > Net >  How to group elements in a list in with respect to another list in Python3?
How to group elements in a list in with respect to another list in Python3?

Time:03-07

I have 2 lists

in1=[1,1,1,2,2,3,4,4,4,5,5]

in2=['a','b','c','d','e','f','g','h','i','j','k']

I want to group the second list based on the same elements in the first list i.e.

Output has to be

out=[['a','b','c'],['d','e'],['f'],['g','h','i'],['j','k']]

Explanation: the first 3 elements of the first list are the same, so I want the first 3 elements of the 2nd list to be grouped together (and so on)

If anyone can help out, it would be great!

~Thanks

CodePudding user response:

Just zip the lists, then itertools to the rescue.

from itertools import groupby

in1 = [1,1,1,2,2,3,4,4,4,5,5]
in2 = ['a','b','c','d','e','f','g','h','i','j','k']    

result = [[c for _, c in g] for _, g in groupby(zip(in1, in2), key=lambda x: x[0])]
print(result)
# [['a', 'b', 'c'], ['d', 'e'], ['f'], ['g', 'h', 'i'], ['j', 'k']]

Non-itertools solution:

in1 = [1,1,1,2,2,3,4,4,4,5,5]
in2 = ['a','b','c','d','e','f','g','h','i','j','k']

result = [[]]
key = in1[0]

for k, v in zip(in1, in2):
    if k != key:
        result.append([])
        key = k
    result[-1].append(v)

print(result)
# [['a', 'b', 'c'], ['d', 'e'], ['f'], ['g', 'h', 'i'], ['j', 'k']]
  • Related