Search code examples
pythonlistarithmetic-expressionselementwise-operations

How to perform element-wise arithmetic operations (e.g. add, subtract, multiply) of two equally shaped lists with arbitrary nestings


I want to perform element-wise mathematical operations (e.g. sum, multiply..) on two Python lists containing numbers or multiple nested lists which may contain again numbers or lists and so on.

The shapes of the two lists are equal when performing the operation. Furthermore, the result should be of the same shape as the two input lists.

A shape may differ in:

  1. length,

  2. width (i.e. number of nestings),

  3. order (e.g. the lists start with a number followed by a nested list, but it may also be that it starts with a nested list, followed by numbers).

The shape changes arbitrarily each time I want to perform a mathematical operation.

How can I perform math operations on arbitrarily shaped lists?

In the past, I've implemented a customized piece of code for each different shape, similar to 1, 2, 3, 4, 5, 6, 7, 8, 9, but I was wondering if there is already a more general solution in the form of a library or reference code for this.

Example 1 (summation):

a = [ 1, 2, 3, 4,[ 5, 6, 7, 8]]
b = [10,20,30,40,[50,60,70,80]]
c = elementwiseSUM(a,b)
c

would result in

[11, 22, 33, 44, [55, 66, 77, 88]]

Example 2 (summation):

d = [ 1,[ 2, 3],  4, [ 5, 6, 7], [[ 8], [ 9, 10]]]
e = [10,[20,30], 40, [50,60,70], [[80], [90,100]]]
f = elementwiseSUM(d,e)
f

would result in

[11, [22, 33], 44, [55, 66, 77], [[88], [99, 110]]]

Example 3 (multiplication):

g = [[4,2],1,1]
h = [[8,3],1,9]
i = elementwiseMUL(g,h)
i

would result in

[[32, 6], 1, 9]

elementwiseSUM() and elementwiseMUL() are placeholders for the library functions which I am looking for.


Solution

  • Here's the solution I came up with

    a = [ 1,  2,  3,  4, [ 5,  6,  7,  8]]
    b = [10, 20, 30, 40, [50, 60, 70, 80]]
    
    def element_wise(a, b, f):
        return [element_wise(i, j, f) if type(i) == list and type(j) == list else f(i, j) for i, j in zip(a, b)]
    
    c = element_wise(a, b, lambda x, y: x + y) # [11, 22, 33, 44, [55, 66, 77, 88]]
    

    so a and b are your lists, and f is the function you want to apply, as you can see I wrote a simple function to add ints

    You can either write your own lambdas or just use the operator module.