Search code examples
pythonlist-comprehension

Multiply all elements in 2D list with formula


Problem

I'm trying to change this 2D list by multiplying all elements with a formula, so this:

lst = [[(-1, -1), (-1, 0), (-1, 1)],
       [(0, -1), (0, 0), (0, 1)],
       [(1, -1), (1, 0), (1, 1)]]

using this formula

(1 / (math.pi * 40.5)) * math.e ** -((x ** 2 + y ** 2) / 40.5)

becomes this

[[-0.007480807217716918, -0.007667817778372781, -0.007480807217716918],
[-0.007667817778372781, 1.5, -0.007667817778372781], 
[-0.007480807217716918, -0.007667817778372781, -0.007480807217716918]]

What I tried:

for i in lst:
    for j in i:
        x = j[0]
        y = j[1]
        if x == y == 0:
            newlst.append(1.5)
        else:
            formula = -(1 / (math.pi * 40.5)) * math.e ** -((x ** 2 + y ** 2) / 40.5)
            newlst.append(formula)

print(newlst)

Output:

[-0.007480807217716918, -0.007667817778372781, -0.007480807217716918, -0.007667817778372781, 1.5, -0.007667817778372781, -0.007480807217716918, -0.007667817778372781, -0.007480807217716918]

The issue is I've tried to solve using list comprehension but failed syntax every time. I also want it to return a 2D list like in lst.


Solution

  • Something like this should work if you want list comprehension

    lst = [(1 / (math.pi * 40.5)) * math.e ** -((x ** 2 + y ** 2) / 40.5) for j in lst for (x,y) in j]
    

    Using a for loop could look like this:

    for i, row in enumerate(lst):
        for j, (x,y) in enumerate(row):
            if x == y == 0:
                val = 1.5
            else:
                val = (1 / (math.pi * 40.5)) * math.e ** -((x ** 2 + y ** 2) / 40.5)
            lst[i][j] = val
    

    output:

    [[0.007480807217716918, 0.007667817778372781, 0.007480807217716918],
     [0.007667817778372781, 1.5, 0.007667817778372781],
     [0.007480807217716918, 0.007667817778372781, 0.007480807217716918]]