Search code examples
pythonpython-3.xlistgrouping

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


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


Solution

  • 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']]