CODEWITHSUNDEEP
pythonnumpymatrix
David
David

Reputation: 481

Convert a matrix of positive integer numbers into a boolean matrix without loops

I'm trying to write code in Python using NumPy. I'm not sure it's possible but here's what I'm trying to do:

I have a 2D matrix a of shape (rows, cols) with positive integer numbers and I want to define a matrix b such that if a[i,j]=x then b[i,j+1]=b[i,j+2]=...=b[i,j+x]=1 (b is initialized to a matrix of zeros).

You can assume that for every j,x: j+x<=cols-1.

For example, if a is:

[0 2 0 0]
[0 2 0 0]
[3 0 0 0]
[2 0 1 0]

Then b should be:

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

Is it possible to do the above in Python with NumPy without using loops?

If it's not possible to do it without loops, is there an efficient way to do it? (rows and cols can be big numbers.)

Upvotes: 3

Views: 161

Answers (2)

kuzand
kuzand

Reputation: 9806

Here's a slightly optimized solution of @finefoot

aa = a.ravel()
b = np.zeros_like(aa)
for i, x in enumerate(aa):
    if x != 0:
        b[i + 1 : i + 1 + x] = 1
b = b.reshape(a.shape)

And here's another solution which is slightly faster but less readable:

from itertools import chain

aa = a.ravel()
b = np.zeros_like(aa)
w = np.nonzero(aa)[0]
ranges = (range(s, e) for s, e in zip(w + 1, w + 1 + aa[w]))
for r in chain.from_iterable(ranges):
    b[r] = 1
b = b.reshape(a.shape)

Gives correct results under assumption that j,x: j+x<=cols-1. Both solutions use a for-loop though, but I don't think that it's possible to do it otherwise.

Upvotes: 0

finefoot
finefoot

Reputation: 11232

If it's not possible to do it without loops, is there an efficient way to do it? (rows and cols can be big numbers.)

I'm sorry, I don't know a NumPy function which would help in your situation, but I think a regular loop and array indexing should be quite fast:

import numpy as np

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

b = np.zeros(a.shape)
for i, x in enumerate(a.flat):
    b.flat[i + 1 : i + 1 + x] = 1

print(b)

Which prints your expected result:

[[0. 0. 1. 1.]
 [0. 0. 1. 1.]
 [0. 1. 1. 1.]
 [0. 1. 1. 1.]]

Upvotes: 1

Related Questions