Как ускорить работу классификатора knn?

Я делаю задание, где нужно написать свой аналог sklearn knn. Вот основной код класса:

 def fit(self, X, y):
        self.X_train = X
        self.scaler = preprocessing.StandardScaler().fit(self.X_train) 
        self.X_train_st = self.scaler.transform(self.X_train)
        self.y_train = y
        self.min_length = np.unique(self.y_train)[-1]+1
        
    def predict(self, X):
        self.X_test = X
        self.X_test_st = self.scaler.transform(self.X_test)
        self.y_pred = np.zeros(self.X_test_st.shape[0],dtype='int64')
        if(self.alg == 'brute'):
            dist = distance.cdist(self.X_test_st, self.X_train_st,  'minkowski', p=2.)
            ind = np.argsort(dist,axis=1)[:,:self.n]
            self.y_pred = np.argmax(np.apply_along_axis(lambda x: np.bincount(x, minlength=self.min_length), 1, self.y_train[ind]), axis=1)
        else:
            self.KDTree = KDTree(self.X_train_st)
            dist, ind = self.KDTree.query(self.X_test_st, k=self.n)
            self.y_pred = np.argmax(np.apply_along_axis(lambda x: np.bincount(x, minlength=self.min_length), 1, self.y_train[ind]), axis=1)
        return self.y_pred.ravel()

На всем известных ирисах работает достаточно быстро по сравнению с sklearn, но я хочу попробовать с fetch_20newsgroups и тут уже все плохо. Сначала я использовала

vectorizer = feature_extraction.text.CountVectorizer(max_df = 0.15, min_df=0.01)
X = vectorizer.fit_transform(data)

Попробовала посмотреть средний скор на кросс-валидации на трех фолдах, brute считается просто нереально долго, по сравнению с тем же sclearn, kdtree немного получше, но все равно очень долго. TfidfVectorizer и cosine метрика помогают немного с точностью(хотя все равно очень низкая, около 15-17%), но что делать со скоростью я не знаю, особенно учитывая, что косинусную метрику использовать на кд дереве нельзя.

UPD: на pairwise_distances и без стандартизации все взлетело и по точности, и по скорости)