Search code examples
pythonfunctools

How to reduce lists in a list


My gut says that reduce (from functools import reduce) would deliver what I'm looking for, yet I can't wrap my head around using it here.

def test_reduce_lists_by_summing_them():
    """
    Sum each item from the first list with the same positional item from the subsequent list. 
    The result is the next first list and appended to the returned result list.
    """
    input_ = [[1, 0, 0], [0, 1, 0], [0, 0, 2], [3, 0, 0]]
    expected_output = [[1, 1, 0], [1, 1, 2], [4, 1, 2]]

    def f(x, y):
        """Return the sum of two positional items"""
        return x + y

    output = []
    for group in zip(input_[0:], input_[1:]):
        # Obvisously not working since the first list is not generated but just taken grp[0]
        output.append(list(map(f, group[0], group[1])))
    assert output == expected_output

Solution

  • What you're looking for is accumulate:

    from itertools import accumulate
    
    # perhaps, convert it to a list explicitly
    accumulate(input_, lambda x, y: list(map(sum, zip(x, y))))
    
    # a more explicit equivalent:
    list(accumulate(input_, lambda x, y: [x_+y_ for x_, y_ in zip(x, y)]))
    

    And if you're willing to use numpy, it would be as simple as:

    np.array(input_).cumsum(axis=0)
    

    UPD Demo:

    In [36]: list(accumulate(input_, lambda x, y: list(map(sum, zip(x, y)))))
    Out[36]: [[1, 0, 0], [1, 1, 0], [1, 1, 2], [4, 1, 2]]
    
    In [37]: list(accumulate(input_, lambda x, y: [x_+y_ for x_, y_ in zip(x, y)]))
    Out[37]: [[1, 0, 0], [1, 1, 0], [1, 1, 2], [4, 1, 2]]
    
    In [38]: np.array(input_).cumsum(axis=0)
    Out[38]: 
    array([[1, 0, 0],
           [1, 1, 0],
           [1, 1, 2],
           [4, 1, 2]])