It turns out, that che chaindes version is the fastest ... so far.

library/NNSearch.py

@@ -3,6 +3,7 @@ import math
 import tensorflow as tf
 import numpy as np
 from sklearn.neighbors import NearestNeighbors
+from library.timing import timing
 
 
 def dist(x,y):
@@ -311,19 +312,36 @@ class Ball:
 
 
 class NNSearch:
-    def __init__(self, nebSize=5):
+    def __init__(self, nebSize=5, timingDict=None):
         self.nebSize = nebSize
         self.neighbourhoods = []
+        self.timingDict = timingDict
 
 
+    def timerStart(self, name):
+        if self.timingDict is not None:
+            if name not in self.timingDict:
+                self.timingDict[name] = timing(name)
+
+            self.timingDict[name].start()
+
+    def timerStop(self, name):
+        if self.timingDict is not None:
+            if name in self.timingDict:
+                self.timingDict[name].stop()
+
     def neighbourhoodOfItem(self, i):
         return self.neighbourhoods[i]
 
 
     def fit(self, X, nebSize=None):
-        self.fit_bruteForce_np(X, nebSize)
+        self.timerStart("NN_fit_all")
+        self.fit_chained(X, nebSize)
+        #self.fit_bruteForce_np(X, nebSize)
+        self.timerStop("NN_fit_all")
        
     def fit_optimized(self, X, nebSize=None):
+        self.timerStart("NN_fit_all")
         if nebSize == None:
             nebSize = self.nebSize
 
@@ -336,28 +354,40 @@ class NNSearch:
         else:
             print("Using chained")
             self.fit_chained(X, nebSize)
+        self.timerStop("NN_fit_all")
 
 
     def fit_bruteForce(self, X, nebSize=None):
         if nebSize == None:
             nebSize = self.nebSize
 
+        self.timerStart("NN_fit_bf_init")
         isGreaterThan = lambda x, y: x[1] > y[1]
         self.neighbourhoods = [MaxHeap(nebSize, isGreaterThan, (i, 0.0)) for i in range(len(X))]
+        self.timerStop("NN_fit_bf_init")
 
+        self.timerStart("NN_fit_bf_loop")
         for (i, x) in enumerate(X):
             nbh = self.neighbourhoods[i]
 
             for (j, y) in enumerate(X[i+1:]):
                 j += i + 1
+                self.timerStart("NN_fit_bf_dist")
                 d = dist(x,y)
+                self.timerStop("NN_fit_bf_dist")
+                self.timerStart("NN_fit_bf_insert")
                 nbh.insert((j,d))
                 self.neighbourhoods[j].insert((i,d))
+                self.timerStop("NN_fit_bf_insert")
 
+            self.timerStart("NN_fit_bf_toList")
             self.neighbourhoods[i] = nbh.toArray(lambda v: v[0])
+            self.timerStop("NN_fit_bf_toList")
+        self.timerStop("NN_fit_bf_loop")
 
 
     def fit_bruteForce_np(self, X, nebSize=None):
+        self.timerStart("NN_fit_bfnp_init")
         numOfPoints = len(X)
         nFeatures = len(X[0])
 
@@ -374,21 +404,31 @@ class NNSearch:
 
         isGreaterThan = lambda x, y: x[1] > y[1]
         self.neighbourhoods = [MaxHeap(nebSize, isGreaterThan, (i, 0.0)) for i in range(len(X))]
+        self.timerStop("NN_fit_bfnp_init")
 
+        self.timerStart("NN_fit_bfnp_loop")
         for (i, x) in enumerate(X):
+            self.timerStart("NN_fit_bfnp_dist")
             distances = distancesTo(x)
+            self.timerStop("NN_fit_bfnp_dist")
             nbh = self.neighbourhoods[i]
 
             for (j, y) in enumerate(X[i+1:]):
                 j += i + 1
                 d = distances[j]
+                self.timerStart("NN_fit_bfnp_insert")
                 nbh.insert((j,d))
                 self.neighbourhoods[j].insert((i,d))
+                self.timerStop("NN_fit_bfnp_insert")
 
+            self.timerStart("NN_fit_bfnp_toList")
             self.neighbourhoods[i] = nbh.toArray(lambda v: v[0])
+            self.timerStop("NN_fit_bfnp_toList")
+        self.timerStop("NN_fit_bfnp_loop")
 
 
     def fit_chained(self, X, nebSize=None):
+        self.timerStart("NN_fit_chained_init")
         if nebSize == None:
             nebSize = self.nebSize
 
@@ -397,10 +437,13 @@ class NNSearch:
 
         neigh = NearestNeighbors(n_neighbors=nebSize)
         neigh.fit(X)
+        self.timerStop("NN_fit_chained_init")
+        self.timerStart("NN_fit_chained_toList")
         self.neighbourhoods = [
                 (neigh.kneighbors([x], nebSize, return_distance=False))[0]
                 for (i, x) in enumerate(X)
                 ]
+        self.timerStop("NN_fit_chained_toList")
 
 
     # ===============================================================