Mejora de la velocidad del algoritmo de agrupamiento codicioso

Estoy tratando de implementar un algoritmo de agrupamiento codicioso muy simple en Python, pero estoy presionado para optimizarlo para la velocidad. El algoritmo tomará una matriz de distancia, buscará la columna con la mayoría de los componentes menos que un corte de distancia predeterminada y almacenará los índices de fila (con componentes menos que el corte) como miembros del grupo. El centroide del cluster es el índice de la columna. Las columnas y filas de cada índice de miembro se eliminan de la matriz de distancias (lo que da como resultado una matriz más pequeña, pero todavía cuadrada,), y el algoritmo itera a través de matrices de distancias sucesivamente más pequeñas hasta que se encuentran todos los grupos. Debido a que cada iteración depende de la última (se forma una nueva matriz de distancia para que no haya miembros superpuestos entre grupos), creo que no puedo evitar un bucle lento para python. He intentado numba (jit) para acelerarlo, pero creo que está volviendo al modo python y por lo tanto no produce ningún aumento de velocidad. Aquí hay dos implementaciones del algoritmo. El primero es más lento que el segundo. Cualquier sugerencia para aceleraciones son bienvenidas. Soy consciente de otros algoritmos de agrupación en clúster implementados en scipy o sklearn (como DBSCAN, kmeans / medoids, etc.), pero estoy muy interesado en usar el actual para mi aplicación. Gracias de antemano por cualquier sugerencia.

Método 1 (más lento):

def cluster(distance_matrix, cutoff=1): indices = np.arange(0, len(distance_matrix)) boolean_distance_matrix = distance_matrix <= cutoff centroids = [] members = [] while boolean_distance_matrix.any(): centroid = np.argmax(np.sum(boolean_distance_matrix, axis=0)) mem_indices = boolean_distance_matrix[:, centroid] mems = indices[mem_indices] boolean_distance_matrix[mems, :] = False boolean_distance_matrix[:, mems] = False centroids.append(centroid) members.append(mems) return members, centroids 

Método 2 (más rápido, pero aún lento para matrices grandes):

Toma como entrada una matriz de adyacencia (dispersa) formada a partir de la implementación de los vecinos más cercanos de sklearn. Esta es la forma más simple y rápida en la que podría pensar para obtener la matriz de distancia relevante para la agrupación. Creo que trabajar con la matriz dispersa también acelera el algoritmo de agrupamiento.

 nbrs = NearestNeighbors(metric='euclidean', radius=1.5, algorithm='kd_tree') nbrs.fit(data) adjacency_matrix = nbrs.radius_neighbors_graph(data) def cluster(adjacency_matrix, gt=1): rows = adjacency_matrix.nonzero()[0] cols = adjacency_matrix.nonzero()[1] members = [] member = np.ones(len(range(gt+1))) centroids = [] appendc = centroids.append appendm = members.append while len(member) > gt: un, coun = np.unique(cols, return_counts=True) centroid = un[np.argmax(coun)] appendc(centroid) member = rows[cols == centroid] appendm(member) cols = cols[np.in1d(rows, member, invert=True)] rows = rows[np.in1d(rows, member, invert=True)] return members, centroids