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

Question

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 ?

Answer 1

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

Answer 2

Without using itertools:

indices = [1, 3, 5]

def groupWith(elements, indices):
    if sum(indices)<len(elements):
        q, r = divmod(len(elements)-sum(indices), indices[-1])
        indices += [indices[-1]] * q + ([r] if r>0 else [])
    elif sum(indices)>len(elements):
        cumsum_indices = [sum(indices[:i+1]) for i in range(len(indices))]
        for j, e in enumerate(cumsum_indices):
            if e>len(elements):
                indices = indices[:j] + [len(elements)-sum(indices[:j])]
                break
    ret = []
    for idx, elem in enumerate(indices):
        start = int(idx>0) * sum(indices[:idx])
        ret += [elements[start:start+elem]]
    return ret
    
print(groupWith(['A','B','C','D','E','F','G','H','I'], indices))
print(groupWith(['A','B','C','D','E','F','G','H','I', 'J','K','L','M','N'], indices))
print(groupWith(['A','B','C','D','E','F','G','H','I', 'J','K','L'], indices))
print(groupWith(['A','B','C','D','E','F','G','H','I', 'J','K','L','M','N','O','P','Q','R','S'], indices))
print(groupWith(['A','B','C'], indices))
print(groupWith(['A','B','C','D','E','F'], indices))

Outputs:

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

Note: this is a bit crude and could very well be optimised.