Edit: this answer refers to v1 of the preprint in the link OP provided. I didn't look closely at the subsequent versions.
As explained in the link you provided, given access to an oracle that prepares each state in the set $\{ |\phi_i\rangle\}$ for $i=1, \dots, M$ and some fixed state $|\psi\rangle$, you can estimate $F_i =|\langle \psi |\phi_i \rangle|^2$ for $i=1, \dots, M$ using many repetitions of the circuit provided in Figure 4 of that preprint. If you read through Section IIIA it will describe how to modify the standard SWAP test for this purpose. The only major change is that you need to add a fourth register containing $m=\lceil \log (\text{len}(A))\rceil$ qubits (3 in your case) and have a unitary $W$ that acts on your third and fourth registers according to
$$ \tag{1}
W|0\rangle^{\otimes n}|i\rangle = |\phi_i\rangle |i\rangle
$$
If you read carefully through the rest of that section it will describe the outputs of this circuit and therefore how it samples approximations of $F_i$ according to their magnitude, thereby sampling the distances of $|\psi\rangle$ to its K nearest neighbors. However, you are probably better off just using your code for a normal SWAP test with randomly sampled $|\phi_i\rangle$, as I will demonstrate below.
An equivalent alternative
I'm going to demonstrate why this is maybe not such a good idea compared to just performing the standard SWAP test. Consider Equations 18 and 22 in the preprint, the first of which describes the probability of measuring the first qubit as "0" and the second describes the probability of measuring the fourth register in "i" conditioned on the first register being "0":
\begin{align}
p(0) &= \frac{1}{2} + \frac{1}{2M} \sum_{j=1}^M F_j \tag{2, 18 in preprint}\\
p(i|0) &= \frac{1 + F_i}{M + \sum_{j=1}^M F_j} \tag{3, 22 in preprint}
\end{align}
Then the joint probability distribution over both registers is just
\begin{align}
p(i, 0) &= p(i|0) p(0) \tag{4}
\\&= \frac{1 + F_i}{M + \sum_{j=1}^M F_j} \cdot \frac{1}{2M}\left(M + \sum_{j=1}^M F_j \right) \tag{5}
\\&= \frac{1 + F_i}{2M} \tag{6}
\end{align}
Now consider a different protocol with joint distribution $p'(i, 0)$ over the same registers. This protocol is very simple: Randomly uniformly sample $i \in \{1, \dots, M\}$, and then perform a standard SWAP test between $|\psi\rangle$ and $|\phi_i\rangle$. Using standard results from the SWAP test we have
\begin{align}
p'(i, 0) &= p'(i) p'(0|i) \tag{7}
\\&= \frac{1}{M} \left(\frac{1}{2} + \frac{1}{2} F_i\right) \tag{8}
\\&= \frac{1 + F_i}{2M} \tag{9}
\\&= p(i, 0) \tag{10}
\end{align}
This demonstrates that we can exactly reproduce the output distribution of their proposed protocol using classical random numbers instead of quantum superposition (and significantly fewer operations and qubits). Of course, their approach uses marginals $p(i|0)$ which we can't directly access here, but we can of course compute these just using postprocessing on $p'(i, 0)$.
Why this algorithm isn't necessary
Intuitively what's going on here is that when they apply $H^{\otimes m}$ to the fourth register that they introduce, they put the register into a uniform superposition over all computational basis states $\{|i\rangle \}$ for $i=1, \dots, M$. Then the $W$ operator combined with a SWAP test is able to slightly bias sampling of states $i$ according to their inner product $F_i$ with $|\psi\rangle$, but the amount of bias is bounded by the fact that $F_i \leq 1$. So your chances of finding the K nearest neighbors can only improve from an initial $1/M$ to $2/M$. But this is exactly how much you can improve the chance of success for a standard SWAP test when you initially draw $i$ with probability $1/M$
On a more intuitive level, the problem here is that they do not actually use the superposition of the fourth register in a quantum way. Compare their circuit in Figure 4 to the circuit for the Deutsch-Jozsa algorithm, for example, in which case an additional layer of Hadamards is applied at the end of the register to take advantage of destructive intereference of the amplitudes. In contrast, here they just use registers 3 and 4 to prepare a uniform superposition over $\{|\phi_i\rangle\}$ before feeding it into the SWAP test.