Compare neighbors in array

Question

I have a a python function for taking in a 2D numpy array and checking if each element is the same as its neighbor elements. I feel like there's a more efficient way to do this but I'm not sure. Here is the code:

import numpy as np

def compare_neighbors(arr):
    '''
    Checks if element (i,j) is different than (i-1,j),(i+1,j),(i,j-1), or 
    (i,j+1).

    --Input--
    arr: (2D np.array) array to compare all elements of

    --Returns--
    comp_arr: (2D bool np.array) bool array with the resulting comparisons.
           True means the original element is the same as its neighbors,
           False means it was different than at least neighbor
    '''
    comp_arr = np.full(arr.shape, False, dtype=bool) #initialize
    arr_height = arr.shape[0]
    arr_width = arr.shape[1]

    for i in range(arr_height): #Row
        for j in range(arr_width): #column

            center = arr[i,j]

            #Check edges
            if i == 0: #left side
                left = arr[i,j]
            else:
                left = arr[i-1, j]

            if i == arr_height - 1: #right side
                right = arr[i,j]
            else:
                right = arr[i+1,j]

            if j == 0: #up
                up = arr[i,j]
            else:
                up = arr[i, j-1]

            if j == arr_width - 1: #down
                down = arr[i,j]
            else:
                down = arr[i, j+1]

            comp_arr[i,j] = len(set([left, right, up, down, center])) == 1
    return comp_arr

If it is helpful, here are the tests I used for testing it:

A = np.array([[1,1],
              [1,1]])

comp_arr_A = compare_neighbors(A)

B = np.array([[2,2],
              [2,2]])

comp_arr_B = compare_neighbors(B)

C = np.array([[1,1,1,1,1,1,1,1,1],
              [1,1,1,1,1,1,1,1,1],
              [1,2,2,2,2,2,2,2,1],
              [1,2,2,1,1,1,2,2,1],
              [1,2,2,2,2,2,2,2,1],
              [1,1,1,1,1,1,1,1,1]])

comp_arr_C = compare_neighbors(C)

D = np.array([[1,1,1],
              [1,2,1],
              [1,1,1]])

comp_arr_D = compare_neighbors(D)

print(A)
print()
print(comp_arr_A)
print()
print(B)
print()
print(comp_arr_B)
print()
print(C)
print()
print(comp_arr_C)
print()
print(D)
print()
print(comp_arr_D)

which returns

[[1 1]
 [1 1]]

[[ True  True]
 [ True  True]]

[[2 2]
 [2 2]]

[[ True  True]
 [ True  True]]

[[1 1 1 1 1 1 1 1 1]
 [1 1 1 1 1 1 1 1 1]
 [1 2 2 2 2 2 2 2 1]
 [1 2 2 1 1 1 2 2 1]
 [1 2 2 2 2 2 2 2 1]
 [1 1 1 1 1 1 1 1 1]]

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

[[1 1 1]
 [1 2 1]
 [1 1 1]]

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

as expected. All I need the function to do is check each element and compare it to its left, right, above, and bellow neighbors. If it is the same as them then the compare_array in that index is True and if it is different than any of its neighbors then False.

Answer 1

You can improve iterating over the array by using np.ndenumerate to get the current coordinates and current item. From the coordinates you can derive the neighbouring elements.
You can also use inverse checking to only set the respective field to False if a neighbor does not match:

from contextlib import suppress

def compare_neighbors(arr):

    comp_arr = np.full(arr.shape, True, dtype=bool)

    for (x, y), item in np.ndenumerate(arr):
        # Check left.
        if x >= 0:
            if arr[x-1, y] != item:
                comp_arr[x, y] = False
                continue

        # Check right.
        with suppress(IndexError):
            if arr[x+1, y] != item:
                comp_arr[x, y] = False
                continue

        # Check top.
        with suppress(IndexError):
            if arr[x, y+1] != item:
                comp_arr[x, y] = False
                continue

        # Check bottom.
        if y >= 0:
            if arr[x, y-1] != item:
                comp_arr[x, y] = False
                continue

    return comp_arr

Answer 2

You can also use the np.roll function if you have a bigger array like in a post by doing:

import numpy as np
def shift_helper(array, shift=0, axis=0):
    # Roll the array by n unity along one axis
    _array = np.roll(_array, shift=shift, axis=axis)

    # Cancel the last/first slice rolled to the first/last slice
        if axis == 0:
            if shift >= 0:
                _array[:1, :, :] = False
            else:
                _array[-1:, :, :] = False
            return _array
        elif axis == 1:
            if shift >= 0:
                _array[:, :1, :] = False
            else:
                _array[:, -1:, :] = False
            return _array

        #Uncomment it for 3D array
        #elif axis == 2:
            #if shift >= 0:
                #_array[:, :, :1] = False
            #else:
                #_array[:, :, -1:] = False
            #return _array

def compare(array, that_value):
    bool_array = np.zeros(array.shape, dtype=bool)
    bool_array[np.where((array == that_value)
                        & (shift_helper(array!=that_value, shift=1, axis=0)#up
                        | shift_helper(array!=that_value, shift=-1, axis=0)#down
                        | shift_helper(array!=that_value, shift=1, axis=1)#left
                        | shift_helper(array!=that_value, shift=-1, axis=1)#right
                        #Uncomment below for 3D array
                        #| shift_helper(array!=that_value, shift=1, axis=2)#front
                        #| shift_helper(array!=that_value, shift=-1, axis=2)#back
                ))] = True
    return bool_array
# Main
C = np.array([[1,1,1,1,1,1,1,1,1],
              [1,1,1,1,1,1,1,1,1],
              [1,2,2,2,2,2,2,2,1],
              [1,2,2,1,1,1,2,2,1],
              [1,2,2,2,2,2,2,2,1],
              [1,1,1,1,1,1,1,1,1]])
print(compare(C, 1))
print(compare(C, 2))

You can choose the value that you want to compare and extend to 3D with this.