CODEWITHSUNDEEP
pythonnumpy
Coffee_Table
Coffee_Table

Reputation: 284

Efficiently create numpy array with decreasing values moving away from center

Let N be an odd positive integer. I want to create a square 2D grid of shape N x N where the center element is N, the surrounding 8-neighborhood has value N-1, the elements around that neighborhood have value N-2 etc., so that finally the "outer shell" of the array has value N//2. I can accomplish this using np.pad in a for loop that iteratively adds "shells" to the array:

def pad_decreasing(N):
    assert N % 2 == 1
    a = np.array([[N]])
    for n in range(N-1, N//2, -1):
        a = np.pad(a, 1, mode='constant', constant_values=n)
    return a

Example output:

In: pad_decreasing(5)
Out: 
array([[3, 3, 3, 3, 3],
       [3, 4, 4, 4, 3],
       [3, 4, 5, 4, 3],
       [3, 4, 4, 4, 3],
       [3, 3, 3, 3, 3]])

Question. Can I accomplish this in a "vectorized" form that doesn't resort to for loops. This code is fairly slow for N large.

Upvotes: 2

Views: 586

Answers (1)

akuiper
akuiper

Reputation: 215057

The distance between the center and any point in the grid can be written as np.maximum(np.abs(row - center), np.abs(col - center)); Then N minus this distance gives the onion you need: 

N = 5

# the x and y coordinate of the center point
center = (N - 1) / 2

# the coordinates of the square array
row, col = np.ogrid[:N, :N]

# N - distance gives the result
N - np.maximum(np.abs(row - center), np.abs(col - center))

# array([[3, 3, 3, 3, 3],
#        [3, 4, 4, 4, 3],
#        [3, 4, 5, 4, 3],
#        [3, 4, 4, 4, 3],
#        [3, 3, 3, 3, 3]])

Upvotes: 4

Related Questions