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[1]
    col = ind % imsize[1]
    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.

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~

1 Like

Thanks for your reply! It helps me a lot.

Hi, great it helps! If it has solved your issue, please help me to set it as the solution so that others could also benefit from our discussions, thank you

If not, then just ignore this pose lol, cheers~

2 Likes

Sure. It is great! Thank you again.