One thing you aren't taking into account is that the computation of the similarity scores is O(n) but the sorting he's doing is n log n which for 1m might dominate especially since it's not necessarily hardware optimized

There's an O(n) algorithm for top k partitioning that could be much much faster than .sort() when you have thousands of elements.

QuickSelect. In C++ its available as std::nth_element in swift I couldn't find it directly but you can implement it in a few lines using .partition as a subroutine