Successive zeroing of columns of a numpy array

Viewed 109

I have an array a of ones and zeroes (it might be rather big)

a = np.array([[1, 0, 0, 1, 0, 0],
              [1, 1, 0, 0, 1, 0],
              [0, 1, 1, 0, 0, 1],
              [0, 0, 0, 1, 1, 1])

in which the "upper" rows are more "important" in the sense that if there is 1 in any column of the i-th row, then all ones in that columns in the following rows must be zeroed.

So, the desired output should be:

array([[1, 0, 0, 1, 0, 0],
       [0, 1, 0, 0, 1, 0],
       [0, 0, 1, 0, 0, 1],
       [0, 0, 0, 0, 0, 0]])

In other words, there should only be single 1 per column.

I'm looking for a more numpy way to do this (i.e. minimising or, better, avoiding the loops).

3 Answers

Your array:

     [[1, 0, 0, 1, 0, 0],
      [1, 1, 0, 0, 1, 0],
      [0, 1, 1, 0, 0, 1],
      [0, 0, 0, 1, 1, 1]]

Transpose it with numpy:

a = np.transpose(your_array)

Now it looks like this:

  [[1, 1, 0, 0],
   [0, 1, 1, 0],
   [0, 0, 1, 0],
   [1, 0, 0, 1],
   [0, 1, 0, 1],
   [0, 0, 1, 1]]

Zero all the non-zero (and "not upper") elements row wise:

 res = np.zeros(a.shape, dtype="int64")
 idx =  np.arange(res.shape[0])
 args = a.astype(bool).argmax(1)
 res[idx, args] = a[idx, args]
 

The output of res is this:

 #### Output
  [[1, 0, 0, 0],
   [0, 1, 0, 0],
   [0, 0, 1, 0],
   [1, 0, 0, 0],
   [0, 1, 0, 0],
   [0, 0, 1, 0]]

Re-transpose your array:

a = np.transpose(res)

  [[1, 0, 0, 1, 0, 0],
   [0, 1, 0, 0, 1, 0],
   [0, 0, 1, 0, 0, 1],
   [0, 0, 0, 0, 0, 0]])

EDIT: Thanks to @The.B for the tip

An alternative solution is to do a forward fill followed by the cumulative sum and then replace all values which are not 1 with 0:

a = np.array([[1, 0, 0, 1, 0, 0],
              [1, 1, 0, 0, 1, 0],
              [0, 1, 1, 0, 0, 1],
              [0, 0, 0, 1, 1, 1]])

ff = np.maximum.accumulate(a, axis=0)
cs = np.cumsum(ff, axis=0)
cs[cs > 1] = 0

Output in cs:

array([[1, 0, 0, 1, 0, 0],
       [0, 1, 0, 0, 1, 0],
       [0, 0, 1, 0, 0, 1],
       [0, 0, 0, 0, 0, 0]])

EDIT

This will do the same thing and should be slightly more efficient:

ff = np.maximum.accumulate(a, axis=0)
ff ^ np.pad(ff, ((1,0), (0,0)))[:-1]

Output:

array([[1, 0, 0, 1, 0, 0],
       [0, 1, 0, 0, 1, 0],
       [0, 0, 1, 0, 0, 1],
       [0, 0, 0, 0, 0, 0]])

And if you want to do the operations in-place to avoid temporary memory allocation:

out = np.zeros((a.shape[0]+1, a.shape[1]), dtype=a.dtype)
np.maximum.accumulate(a, axis=0, out=out[1:])
out[:-1] ^ out[1:]

Output:

array([[1, 0, 0, 1, 0, 0],
       [0, 1, 0, 0, 1, 0],
       [0, 0, 1, 0, 0, 1],
       [0, 0, 0, 0, 0, 0]])

You can traverse through each column of array and check if it is the first one -

If Not: Make it 0

for col in a.T:
  f=0
  for x in col:
    if(x==1 and f==0):
      f=1
    else:
      x=0
Related