Search code examples
pythonlistfunctionindexingpython-itertools

Summing lists of different lengths in Python beginning at a specified index using a function


How can I combine multiple lists starting at a specified index within one of the lists, while also summing the values at the respective indexes?

If I had 3 lists:

a = [1, 2, 3, 4, 5, 0, 0, 0, 0]
b = [2, 3, 4, 5, 6]
c = [5, 2]

How could I make it so that I could insert lists b and c into any position within a, and the respective indexes would sum?

For example, I could insert list b starting from the first index of a, and also insert list c starting from the fifth index of a. The output of this would look something like:

NewList = [1, 4, 6, 8, 10, *11*, 2, 0, 0]

You can see that when adding lists b and c at the aforementioned indexes of a, there was an overlap at the fifth index, which I would like to do.

I have tried using izip_longest from itertools in combination with a slicing method:

result = [sum(n) for n in izip_longest(a[2:], b, fillvalue=0)]  

But this produces:

result = [5, 7, 9, 5, 6, 0, 0]

This cuts off the zero and one index in list a, which I do not want. However, I want to use izip_longest because I can sum lists with different lengths.

I would like to create a function that does this with an arbitrary amount of lists, so a could be a list with len = 1000 and I could insert any amount of different lists with different lengths wherever I want into a, as many times as I want and the index values would sum.


Solution

  • Here is a function that should do what you're asking for:

    def combine(lista, listb, index_in_list_a_to_start_insert):
        # left pad list b with zeros
        newb = [0]*index_in_list_a_to_start_insert + listb
    
        # right pad shorter list
        max_len = max(len(lista), len(newb))
        newa = lista if len(lista) >= max_len else lista + [0]*(max_len-len(lista))
        newb = newb if len(newb) >= max_len else newb + [0]*(max_len-len(newb))
    
        # sum element-wise            
        return [a + b for a, b in zip(newa,newb)]
    

    You can run your example by chaining this function as so:

    a = [1, 2, 3, 4, 5, 0, 0, 0, 0]
    b = [2, 3, 4, 5, 6]
    c = [5, 2]
    combine(combine(a, b, 1), c, 5)
    

    Output:

     [1, 4, 6, 8, 10, 11, 2, 0, 0]
    

    You'd probably also want to add some error checking to make sure that the specified index is in-bounds.