Can you use if-else statement inside braces for an array in Python?

Question

I'm wondering if there is a nice way to use an if-else statement inside braces of an array in Python to assign values. What I would like is something like:

A = #some 2D array of length m by n, already initialized
A = np.float64(A)
val = someValue #any number, pick a number

A = [[val for j in range(n) if A[i][j] < val, else A[i][j]=A[i][j]] for i in range(m)]

Is there a nice way to do this? Alternatively, if numpy has a faster way to compute this that would be equally as good, if not better.

The longer way to do what I am trying to achieve would be something like

for i in range(m):
    for j in range(n):
        if A[i][j] < val:
            A[i][j] = val

The desired output is to set any values below a threshold to that threshold. I can do simpler if-statements with a 1D array such as

myArray = [otherArray[i] for i in range(theRange) if otherArray[i]>=value and otherArray[i]<=anotherValue]

This 1D example is not what I want. It's just an example of the type of coding block I'm looking for. It seems to be quicker at processing against the traditional if-else statements.

Answer 1

With numpy arrays we try avoid iteration (list comprehension). Sometimes it is needed, but in this case it is not:

In [403]: A=np.arange(16).reshape(4,4)    
In [404]: A1=A.astype(np.float64)    # better syntax for converting to float

In [405]: A1
Out[405]: 
array([[  0.,   1.,   2.,   3.],
       [  4.,   5.,   6.,   7.],
       [  8.,   9.,  10.,  11.],
       [ 12.,  13.,  14.,  15.]])

A boolean array that shows where the test is True/False:

In [406]: A1<5 
Out[406]: 
array([[ True,  True,  True,  True],
       [ True, False, False, False],
       [False, False, False, False],
       [False, False, False, False]], dtype=bool)

We can index with such a mask:

In [407]: A1[A1<5]=5

In [408]: A1
Out[408]: 
array([[  5.,   5.,   5.,   5.],
       [  5.,   5.,   6.,   7.],
       [  8.,   9.,  10.,  11.],
       [ 12.,  13.,  14.,  15.]])

np.where (and np.nonzero ) return indices where the condition is True; where has a version that operates like the ternary operator (on each element):

In [410]: np.where(A<5,5,A)
Out[410]: 
array([[ 5,  5,  5,  5],
       [ 5,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

We can also clip with np.maximum :

In [411]: np.maximum(A,5)
Out[411]: 
array([[ 5,  5,  5,  5],
       [ 5,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

In [417]: A.clip(5,None)
Out[417]: 
array([[ 5,  5,  5,  5],
       [ 5,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

Answer 2

Python's one-line ternary operator syntax looks like this

variable = a if CONDITION else b

You can place this inside a list comprehension as well. It's not clear what val is in your example, but I'm assuming it's a value you specified beforehand.

val = 2
A = [[val if A[i][j] < val else A[i][j] for j in range(n)] for i in range(m)]