Hi @noUmbrella , then that only needs some additional adjustments:
import torch
import numpy as np
# This is your additional adjustments
P = [2,2]
m = 1
n = 1
N = 3
# DEFINE ARRAY
Q = [[1,2,3,4,5],
[5,6,9,2,1],
[2,5,4,3,6],
[3,6,2,5,1]]
Qnp = np.array(Q)
Qten = torch.tensor(Qnp)
# Find the Top-K (aka. N in your case)
# This is your additional adjustments
out = torch.topk(Qten[P[0]-m:P[0]+m+1,P[1]-n:P[1]+n+1].flatten(), N)
start = np.array([P[0]-m+1, P[1]-n+1, P[0]+m+2, P[1]+n+2])
# get the value and indexes of the desired top-K or N
val, ind = out
# reformat the flatten tensor to your 2d coordinates
def find_coord(ind, start, imsize=[5,5]):
row = ind // imsize[1]
col = ind % imsize[1]
print(row, col)
row = row + start[0]
col = col + start[1]
return np.array([row , col ]) # because you want the array to start from 1 instead of zero
# get the 2d coordinates
all_inds = [find_coord(idx, start, imsize=start[2:]-start[:2]) for idx in ind] # use + 1 to get your index from 1
print('All your desired indexes: ', all_inds)
print('All your desired top-{} values: {}'.format(N, val))
'''
All your desired indexes: [array([2., 3.], dtype=float32), array([2., 2.], dtype=float32), array([4., 2.], dtype=float32)]
All your desired top-3 values: tensor([9, 6, 6], dtype=torch.int32)
'''
Besides that, they are pretty much the same with my first comment, hope it helps, cheers~