CODEWITHSUNDEEP
numpy
durden2.0
durden2.0

Reputation: 9542

Fastest way to compare neighboring elements in multi-dimensional numpy array

What is the fastest way to compare neighboring elements in a 3-dimensional array?

Assume I have a numpy array of (4,4,4). I want to loop in the k-direction and compare elements in pairs. So, compare all neighboring elements and assign the lowest index if they are not equal. Essentially this:

if array([0, 0, 0)] != array[(0, 0, 1)]:
    array[(0, 0, 0)] = 111

Thus, the comparisons would be:

(0, 0, 0) and (0, 0, 1)
(0, 0, 1) and (0, 0, 2)
(0, 0, 2) and (0, 0, 3)
(0, 0, 3) and (0, 0, 4)
... for all i and j ...

However, I want to do this for every i and j in the array and writing a standard Python for loop for this on huge arrays with millions of cells is incredibly slow. Is there a more 'standard' numpy way to do this without the explicit for loop?

Maybe there's some trick using the slicing step (i.e array[:,:,::2], array[:,:,1::2])?

Upvotes: 1

Views: 926

Answers (1)

eickenberg
eickenberg

Reputation: 14377

Try np.diff.

import numpy as np
a = np.arange(9).reshape(3, 3)
A = np.array([a, a, a + 1]).T

same_with_neighbor_on_last_axis = np.diff(A, axis=-1) == 0

print A
print same_with_neighbor_on_last_axis

A is constructed to have 2 consecutive equal entries along the third axis,

>>>print A
array([[[0, 0, 1],
        [3, 3, 4],
        [6, 6, 7]],

       [[1, 1, 2],
        [4, 4, 5],
        [7, 7, 8]],

       [[2, 2, 3],
        [5, 5, 6],
        [8, 8, 9]]])

The output vector then yields

>>>print same_with_neighbor_on_last_axis
[[[ True False]
  [ True False]
  [ True False]]

 [[ True False]
  [ True False]
  [ True False]]

 [[ True False]
  [ True False]
  [ True False]]]

Using the axis keyword, you can choose whichever axis you need to do this operation on. If it is all of them, you can use a loop. np.diff does not much else than the following

np.diff(A, axis=-1) == A[..., 1:] - A[..., :-1]

Upvotes: 4

Related Questions