/** * Cuts a tree into several groups by specifying the desired number. * @param k the number of clusters. * @return the cluster label of each sample. */ public int[] partition(int k) { int n = merge.length + 1; int[] membership = new int[n]; IntHeapSelect heap = new IntHeapSelect(k); for (int i = 2; i <= k; i++) { heap.add(merge[n - i][0]); heap.add(merge[n - i][1]); } for (int i = 0; i < k; i++) { bfs(membership, heap.get(i), i); } return membership; }