# Ask for an efficient implementation (1)

Hello, everyone! I have a matrix Q whose size is W×H. For a specified point P in Q whose coordinate is (x,y), I want to get the coordinates of the first N points whose value rank top-N from the range of (x-m, y-n) to (x+m,y+n). How can I address this problem efficiently using PyTorch.

For instance,
Q =
[1,2,3,4,5;
5,6,9,2,1;
2,5,4,3,6;
3,6,2,5,1;
2,4,8,6,2] and the size of Q is 5×5.
When the coordinate of P is (3,3) and m=n=1 and N=3, the N coordinates I want to get are (2,3), (2,2), and (4,2), whose value are 9, 6, 6, respectively.

Hi, for confirmation, you want the output of N points, which is 3, and you also want to get the point coordinates, am I correct?

``````import torch
import numpy as np
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)
out = torch.topk(Qten.flatten(), N)
# 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, imsize=[5,5]):
row = ind // imsize
col = ind % imsize
return row+ 1, col+1 # because you want the array to start from 1 instead of zero

# get the 2d coordinates

all_inds = [find_coord(idx, imsize=Qnp.shape) for idx in ind]
print('All your desired indexes: ', all_inds)
print('All your desired top-{} values: {}'.format(N, val))
``````

Hope it helps, cheers~

Thanks for your reply and help. I am sorry that I may not describe the problem clearly. In fact, the “m” and “n” in my problem are the most important ones. I want to find the coordinates of the top N points surrounding the P, that is the range of (x-m, y-n) to (x+m,y+n), neither from the whole region of Q.

``````import torch
import numpy as np

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)
out = torch.topk(Qten[P-m:P+m+1,P-n:P+n+1].flatten(), N)
start = np.array([P-m+1, P-n+1, P+m+2, P+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
col = ind % imsize
print(row, col)
row = row + start
col = col + start
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~

1 Like