# Is there a dilated k-nearest neighbour solution available fast execution?

**URL:** <https://discuss.pytorch.org/t/is-there-a-dilated-k-nearest-neighbour-solution-available-fast-execution/197257>\
**Category:** Uncategorized\
**Created:** [February 18, 2024, 10:36am UTC](https://discuss.pytorch.org/t/is-there-a-dilated-k-nearest-neighbour-solution-available-fast-execution/197257 "2024-02-18T10:36:25Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Adex](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/adex/32/11880_2.png) [@Adex](https://discuss.pytorch.org/u/Adex)\
**Post date:** [February 18, 2024, 10:36am UTC](https://discuss.pytorch.org/t/is-there-a-dilated-k-nearest-neighbour-solution-available-fast-execution/197257/1 "2024-02-18T10:36:25Z")

</div>

I am implementing the dilated k Neares Neighbour algorithm. The algorithm unfortunately has nested loops. The presence of loops severely hampers the execution speed.

```
import torch
dilation=3
nbd_size=5
knn_key=torch.randint([0,30,(64,12,198,100)])

dilated_keys=torch.zeros([knn_key.shape[0],knn_key.shape[1],knn_key.shape[2],nbd_size])

for i in range(knn_key.shape[0]):
    for j in range(knn_key.shape[1]):
        for k in range(knn_key.shape[2]):
            list_indices=[]
            while (len(list_indices))<nbd_size:
                for l in range(knn_key.shape[3]):
                    if knn_key[i][j][k][l]%dilation==k%dilation:
                        list_indices.append(knn_key[i][j][k][l])
                        if (len(list_indices))>=nbd_size:
                            break
            list_indices_tensor=torch.tensor(list_indices)
            dilated_keys[i][j][k]=list_indices_tensor

```

The variable `knn_key` stores the `100` nearest neighbours among the 1000 data points originally available. The `dilated_keys` stores the `nbd_size=5` selected indices of the neighbours that are used after applying dialation filter. Any help to use broadcasting solution to remove the three nested loops will be highly helpful.

---

<div class="post-metadata">

**Author:** ![Adex](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/adex/32/11880_2.png) [@Adex](https://discuss.pytorch.org/u/Adex)\
**Post date:** [February 26, 2024, 10:13am UTC](https://discuss.pytorch.org/t/is-there-a-dilated-k-nearest-neighbour-solution-available-fast-execution/197257/2 "2024-02-26T10:13:41Z")

</div>

Looks like the inner loop can be reduced as follows,

```auto
import torch

dilation = 3
nbd_size = 5
knn_key = torch.randint(0, 30, (64, 12, 198, 100))

dilated_keys = torch.zeros((knn_key.shape[0], knn_key.shape[1], knn_key.shape[2], nbd_size), dtype=torch.int64)

for i in range(knn_key.shape[0]):
    for j in range(knn_key.shape[1]):
        for k in range(knn_key.shape[2]):
            key = knn_key[i, j, k]
            indices = torch.nonzero(key % dilation == k % dilation).squeeze()
            selected_indices = indices[:nbd_size]
            dilated_keys[i, j, k] = key[selected_indice

```

However three nested loops are still there. Any help will be highly helpful. [[1]](https://stackoverflow.com/questions/78015382/is-there-a-dilated-k-nearest-neighbour-solution-available-fast-execution/78019217#78019217)