python – How can I speed up array iteration?

Question:

There is an image 1752×1264 pixels. Each point contains a False or True value.

Now, to get an array, we execute the following code:

for i in range(array.shape[0]):
      for j in range(array.shape[1]):
              if array[i][j] == True:
                   rect.append([[j,i]])

On average, it takes 2-3 seconds.

How can you speed up its execution?

Answer:

The easiest way is to use the np.nonzero() function.

source matrix example:

In [29]: a = np.random.choice([True, False], size=(5, 3))

In [30]: a
Out[30]:
array([[False,  True, False],
       [False, False, False],
       [False,  True, False],
       [False, False,  True],
       [False, False,  True]])

solution:

r_idx, c_idx = np.nonzero(a)
res = np.column_stack((c_idx, r_idx))

result:

In [33]: res
Out[33]:
array([[1, 0],
       [1, 2],
       [2, 3],
       [2, 4]], dtype=int64)

execution time measurements:

In [46]: a = np.random.choice([True, False], size=(1752, 1264))

In [47]: %%timeit
    ...: res = np.argwhere(a == True)[:, ::-1]
    ...:
    ...:
44.6 ms ± 539 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [48]: %%timeit
    ...: r_idx, c_idx = np.nonzero(a)
    ...: res = np.column_stack((c_idx, r_idx))
    ...:
    ...:
42.7 ms ± 255 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

UPDATE: To get a 3D array like in the question:

res = res.reshape(res.shape[0], 1, -1)
Scroll to Top