python-3.xlistpython-itertoolsmore-itertoolsitertools-groupby

How to group or batch the list of elements based on indices sublist?


How to group or batch the list of elements based on indices sublist ? Below are the possible instances should be considered. Hope that helps better.

1st Instance : Length of elements == sum of all indices elements

indices = [1, 3, 5]
elements = ['A','B','C','D','E','F','G','H','I']

Output:
[['A'], ['B','C','D'],['E','F','G','H','I']]

2nd Instance : Length of elements > sum of all indices elements

Should consider max element of indices and then group accordingly elements list

indices = [1, 3, 5]
elements = ['A','B','C','D','E','F','G','H','I', 'J','K','L','M','N']

Output:
[['A'], ['B','C','D'],['E','F','G','H','I'],['J','K','L','M','N']]

3rd Instance : Length of elements > sum of all indices elements

Should consider max element of indices and then group accordingly elements list

indices = [1, 3, 5]
elements = ['A','B','C','D','E','F','G','H','I', 'J','K','L']

Output:
[['A'], ['B','C','D'],['E','F','G','H','I'],['J','K','L']]

4th Instance : Length of elements > sum of all indices elements

Should consider max element of indices and then group accordingly elements list

indices = [1, 3, 5]
elements = ['A','B','C','D','E','F','G','H','I', 'J','K','L','M','N','O','P','Q','R','S']

Output:
[['A'], ['B','C','D'],['E','F','G','H','I'],['J','K','L','M','N'], ['O','P','Q','R','S']]

5th Instance : Length of elements < sum of all indices elements

indices = [1, 3, 5]
elements = ['A','B','C']

Output:
[['A'], ['B','C']]

6th Instance : Length of elements < sum of all indices elements

indices = [1, 3, 5]
elements = ['A','B','C','D','E','F']

Output:
[['A'], ['B','C','D'],['E','F']]

How can I achieve this in Python ?


Solution

  • You can use itertools.islice to slice the elements:

    from itertools import islice
    
    indices = [1, 3, 5]
    elements = ['A','B','C','D','E','F','G','H','I', 'J','K','L','M','N','O','P','Q','R','S']
    
    def get_cnt(indices):
        if len(indices) == 0:
            return
        yield from indices
        while True:
            yield indices[-1]
    
    out, it, g = [], iter(elements), get_cnt(indices)
    while slice_:=list(islice(it, next(g))):
        out.append(slice_)
    
    print(out)
    

    Prints:

    [['A'], ['B', 'C', 'D'], ['E', 'F', 'G', 'H', 'I'], ['J', 'K', 'L', 'M', 'N'], ['O', 'P', 'Q', 'R', 'S']]