Summing every 3 elements in list aside from 0-- Python

Question

I am having trouble with the following problem (working in Python):

With any given list, the goal is to sum all elements that are nonzero with 3 being the maximum number of elements summed. After summing the elements, the sum would replace the first element and the other two elements (out of the three that were summed) are replaced with 0. We will leave all elements that are 0 in place.

For example:

Input:   a = [2, 5, 0, 7, 6, 5, 0, 0, 6, 7, 8]
Output:  a = [7, 0, 0,18, 0, 0, 0, 0,21, 0, 0]

Another example--

Input: a = [0,6,0,0,7,9,8,9,8,0,0,2]
Desired output-- a = [0,6,0,0,24,0,0,17,0,0,0,2]

Answer 1

def foo(arr, n=3):
    ans = []
    i = 0
    while i < len(arr):
        if arr[i] == 0:
            ans.append(arr[i])
            i += 1
        else:
            ans.append(sum(arr[i:(i + n)]))
            ans.extend([0] * (n - 1))
            i += n
    return ans[:len(arr)]

foo([2, 5, 0, 7, 6, 5, 0, 0, 6, 7, 8])
# [7, 0, 0, 18, 0, 0, 0, 0, 21, 0, 0]

foo([0, 6, 0, 0, 7, 9, 8, 9, 8, 0, 0, 2])
# [0, 6, 0, 0, 24, 0, 0, 17, 0, 0, 0, 2]

Answer 2

Solution that also handles the corner case [1, 0, 2, 0] -> [1, 0, 2, 0]

Code

def func(arr, n = 3):
    # Group elements into non-zero sublists with of length at most 3
    # Or One element sublist with value 0
    group = []
    for i in arr:
        if not group or i == 0 or group[-1][-1] == 0 or len(group[-1])==n:
            group.append([i])
        else:
            group[-1].append(i)
    
    # Flatten each group by summing as first element with remaining 0
    return [sum(sublist) if i == 0 else 0 for sublist in group for i, x in enumerate(sublist)]

Tests

print(func([2, 5, 0, 7, 6, 5, 0 , 0, 6, 7, 8])) 
#Output: [7, 0, 0, 18, 0, 0, 0, 0, 21, 0, 0]

print(func([0,6,0,0,7,9,8,9,8,0,0,2]))  
# Output: [0, 6, 0, 0, 24, 0, 0, 17, 0, 0, 0, 2]

print(func([1, 0, 2, 0]))
# Output: [1, 0, 2, 0]