numpy conditional format of elements

Question

I have a 3D numpy array and want to change a particular element based on a conditional test of another element. (The applications is to change the 'alpha' of a RGBA image array to play with the transparency in a 3D pyqtgraph image - should ideally be pretty fast).

a= np.ones((2,4,5),dtype=np.int) #create a ones array
a[0,0,0] = 3 #change a few values
a[0,2,0] = 3
a[1,2,0] = 3
print(a)
>>>[[[3 1 1 1 1]
  [1 1 1 1 1]
  [3 1 1 1 1]]

 [[1 1 1 1 1]
  [1 1 1 1 1]
  [3 1 1 1 1]]]

Now I want to conditionally test the first element of the lowest dimension(??) and then change the last element based on the result.

if a[0,:,:] > 1:   #this does not work - for example only - if first element > 1
    a[3,:,:]=255   #then change the last element in the same dimension to 255

print(a) #this is my idealized result
>>>[[[3 1 1 1 255]
  [1 1 1 1 1]
  [3 1 1 1 255]]

 [[1 1 1 1 1]
  [1 1 1 1 1]
  [3 1 1 1 255]]]

Answer 1

This seems to do the trick:

mask = a[:,:,0] > 1

a[:,:,4][mask] = 255

So the indexing just needed to be a little different and then it's just standard practice of applying a mask.

edit @Ophion showed this is much better written as:

a[mask,:,-1] = 255