Ask for an efficient implementation (1)

username: noUmbrella
cooked: <p>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.</p>
<p>For instance,<br>
Q =<br>
[1,2,3,4,5;<br>
5,6,9,2,1;<br>
2,5,4,3,6;<br>
3,6,2,5,1;<br>
2,4,8,6,2] and the size of Q is 5×5.<br>
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.</p>

username: BTNug
cooked: <p>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?</p>
<pre><code class="lang-auto">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))
</code></pre>
<p>Hope it helps, cheers~</p>

username: noUmbrella
cooked: <p>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.</p>

username: BTNug
cooked: <p>Hi <a class="mention" href="/u/noumbrella">@noUmbrella</a> , then that only needs some additional adjustments:</p>
<pre><code class="lang-auto">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)
'''
</code></pre>
<p>Besides that, they are pretty much the same with my first comment, hope it helps, cheers~</p>

username: noUmbrella
cooked: <p>Thanks for your reply! It helps me a lot.</p>

username: BTNug
cooked: <p>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</p>
<p>If not, then just ignore this pose lol, cheers~</p>

username: noUmbrella
cooked: <p>Sure. It is great! Thank you again.</p>