Search code examples
pythonlistelementwise-operations

Apply condition in for loop to each element in a list of lists


I am trying to see if each value in array is less then 0 then output 0, else output the number itself. The output must have the same dimensions as the input. Currently the input is a 3 by 4 but output is a list. How do I get the output size the same (3 by 4 array)?

input = [[1,2,3,4],[4,5,-6,-7],[8,9,10,11]]
output= []
for i in input:
    if i < 0:
        value = 0
        output.append(value)
    else:
        value= i
        output.append(value)

Solution

  • Python's NumPy arrays are a much more efficient way of dealing with matrices, which are what you have if all of your nested lists are the same size. If I understood your question, your example can be simplified down to a single line:

    import numpy as np
    inp = np.array([[-1,2,3,4],[4,5,-6,7],[8,9,-10,11]])
    print (inp)
    #[[ -1   2   3   4]
    # [  4   5  -6   7]
    # [  8   9 -10  11]]
    inp[inp < 0] = 0 
    print (inp)
    # [[ 0  2  3  4]
    # [ 4  5  0  7]
    # [ 8  9  0 11]]