@graphty/algorithms 1.1.0 → 1.3.1

package/src/clustering/hierarchical.ts

@@ -6,6 +6,41 @@
  * structure in graphs.
  */
 
+import type {Graph} from "../core/graph.js";
+
+/**
+ * Convert Graph to adjacency set representation for clustering
+ */
+function graphToAdjacencySet(graph: Graph): Map<string, Set<string>> {
+    const adjacency = new Map<string, Set<string>>();
+
+    // Initialize all nodes
+    for (const node of graph.nodes()) {
+        adjacency.set(String(node.id), new Set());
+    }
+
+    // Add edges
+    for (const edge of graph.edges()) {
+        const source = String(edge.source);
+        const target = String(edge.target);
+
+        const sourceSet = adjacency.get(source);
+        if (sourceSet) {
+            sourceSet.add(target);
+        }
+
+        // For undirected graphs, add reverse edge
+        if (!graph.isDirected) {
+            const targetSet = adjacency.get(target);
+            if (targetSet) {
+                targetSet.add(source);
+            }
+        }
+    }
+
+    return adjacency;
+}
+
 export interface ClusterNode<T> {
     id: string;
     members: Set<T>;
@@ -49,11 +84,7 @@ function computeGraphDistance<T>(
             if (neighbors) {
                 for (const neighbor of neighbors) {
                     if (!dist.has(neighbor)) {
-                        const nodeDistance = dist.get(node);
-                        if (nodeDistance !== undefined) {
-                            dist.set(neighbor, nodeDistance + 1);
-                        }
-
+                        dist.set(neighbor, (dist.get(node) ?? 0) + 1);
                         queue.push(neighbor);
                     }
                 }
@@ -65,55 +96,57 @@ function computeGraphDistance<T>(
 }
 
 /**
- * Compute distance between clusters based on linkage method
+ * Compute distance between two clusters based on linkage method
  */
 function clusterDistance<T>(
-    cluster1: Set<T>,
-    cluster2: Set<T>,
+    cluster1: ClusterNode<T>,
+    cluster2: ClusterNode<T>,
     distances: Map<T, Map<T, number>>,
-    method: LinkageMethod,
+    linkage: LinkageMethod,
 ): number {
-    if (cluster1.size === 0 || cluster2.size === 0) {
-        return Infinity;
-    }
-
     const allDistances: number[] = [];
 
-    for (const node1 of cluster1) {
-        for (const node2 of cluster2) {
-            const dist = distances.get(node1)?.get(node2) ?? Infinity;
-            allDistances.push(dist);
+    for (const node1 of cluster1.members) {
+        const node1Distances = distances.get(node1);
+        if (!node1Distances) {
+            continue;
+        }
+
+        for (const node2 of cluster2.members) {
+            const dist = node1Distances.get(node2);
+            if (dist !== undefined) {
+                allDistances.push(dist);
+            }
         }
     }
 
-    switch (method) {
+    if (allDistances.length === 0) {
+        return Infinity;
+    }
+
+    switch (linkage) {
         case "single":
             return Math.min(... allDistances);
         case "complete":
             return Math.max(... allDistances);
         case "average":
             return allDistances.reduce((a, b) => a + b, 0) / allDistances.length;
-        case "ward":
-            // Ward's method minimizes within-cluster variance
-            // For simplicity, we use average here
-            return allDistances.reduce((a, b) => a + b, 0) / allDistances.length;
+        case "ward": {
+            // Simplified Ward's method - minimize within-cluster variance
+            const size1 = cluster1.members.size;
+            const size2 = cluster2.members.size;
+            const avgDist = allDistances.reduce((a, b) => a + b, 0) / allDistances.length;
+            return avgDist * ((size1 * size2) / (size1 + size2));
+        }
         default:
             return Math.min(... allDistances);
     }
 }
 
 /**
- * Agglomerative hierarchical clustering
- * Builds clusters bottom-up by merging closest pairs
- *
- * @param graph - Undirected graph
- * @param linkage - Linkage method for cluster distance
- * @returns Hierarchical clustering result
- *
- * Time Complexity: O(n³) naive, O(n² log n) with heap
- * Space Complexity: O(n²)
+ * Internal implementation of agglomerative hierarchical clustering
  */
-export function hierarchicalClustering<T>(
+function hierarchicalClusteringImpl<T>(
     graph: Map<T, Set<T>>,
     linkage: LinkageMethod = "single",
 ): HierarchicalClusteringResult<T> {
@@ -148,180 +181,134 @@ export function hierarchicalClustering<T>(
 
     // Initialize cluster distances
     for (let i = 0; i < clusters.length; i++) {
-        const clusterI = clusters[i];
-        if (!clusterI) {
+        const cluster1 = clusters[i];
+        if (!cluster1) {
             continue;
         }
 
-        clusterDistances.set(clusterI.id, new Map());
+        const distMap = new Map<string, number>();
         for (let j = i + 1; j < clusters.length; j++) {
-            const clusterJ = clusters[j];
-            if (!clusterJ) {
+            const cluster2 = clusters[j];
+            if (!cluster2) {
                 continue;
             }
 
-            const dist = clusterDistance(
-                clusterI.members,
-                clusterJ.members,
-                distances,
-                linkage,
-            );
-            const clusterIDistances = clusterDistances.get(clusterI.id);
-            if (clusterIDistances) {
-                clusterIDistances.set(clusterJ.id, dist);
-            }
+            const dist = clusterDistance(cluster1, cluster2, distances, linkage);
+            distMap.set(cluster2.id, dist);
         }
+        clusterDistances.set(cluster1.id, distMap);
     }
 
+    // Active clusters
+    const activeClusters = new Set(clusters);
+
     // Merge clusters until only one remains
-    while (clusters.length > 1) {
-    // Find closest pair of clusters
+    while (activeClusters.size > 1) {
+        // Find closest pair of clusters
         let minDist = Infinity;
-        let merge1 = 0;
-        let merge2 = 0;
-
-        for (let i = 0; i < clusters.length; i++) {
-            for (let j = i + 1; j < clusters.length; j++) {
-                const clusterI = clusters[i];
-                const clusterJ = clusters[j];
-                if (!clusterI || !clusterJ) {
+        let mergeCluster1: ClusterNode<T> | null = null;
+        let mergeCluster2: ClusterNode<T> | null = null;
+
+        for (const cluster1 of activeClusters) {
+            const distMap = clusterDistances.get(cluster1.id);
+            if (!distMap) {
+                continue;
+            }
+
+            for (const cluster2 of activeClusters) {
+                if (cluster1.id >= cluster2.id) {
                     continue;
                 }
 
-                const dist = clusterDistances.get(clusterI.id)?.get(clusterJ.id) ??
-                    clusterDistances.get(clusterJ.id)?.get(clusterI.id) ??
+                const dist = distMap.get(cluster2.id) ??
+                    clusterDistances.get(cluster2.id)?.get(cluster1.id) ??
                     Infinity;
 
                 if (dist < minDist) {
                     minDist = dist;
-                    merge1 = i;
-                    merge2 = j;
+                    mergeCluster1 = cluster1;
+                    mergeCluster2 = cluster2;
                 }
             }
         }
 
-        // Check if we found a valid merge
-        if (minDist === Infinity || merge1 === merge2) {
-            // No more clusters can be merged (disconnected components)
+        if (!mergeCluster1 || !mergeCluster2 || minDist === Infinity) {
+            // No valid merge found - create forest
+            const trees = Array.from(activeClusters);
+            const forestRoot: ClusterNode<T> = {
+                id: `forest-${String(clusterCount++)}`,
+                members: new Set(nodes),
+                distance: Infinity,
+                height: Math.max(... trees.map((t) => t.height)) + 1,
+                trees,
+            };
+            dendrogram.push(forestRoot);
+            activeClusters.clear();
+            activeClusters.add(forestRoot);
             break;
         }
 
         // Create new cluster
-        const cluster1 = clusters[merge1];
-        const cluster2 = clusters[merge2];
-        if (!cluster1 || !cluster2) {
-            continue;
-        }
-
         const newCluster: ClusterNode<T> = {
             id: `cluster-${String(clusterCount++)}`,
-            members: new Set([... cluster1.members, ... cluster2.members]),
-            left: cluster1,
-            right: cluster2,
+            members: new Set([... mergeCluster1.members, ... mergeCluster2.members]),
+            left: mergeCluster1,
+            right: mergeCluster2,
             distance: minDist,
-            height: Math.max(cluster1.height, cluster2.height) + 1,
+            height: Math.max(mergeCluster1.height, mergeCluster2.height) + 1,
         };
 
         dendrogram.push(newCluster);
 
-        // Update cluster distances for new cluster
-        clusterDistances.set(newCluster.id, new Map());
+        // Update distances to new cluster
+        const newDistMap = new Map<string, number>();
+        for (const cluster of activeClusters) {
+            if (cluster.id === mergeCluster1.id || cluster.id === mergeCluster2.id) {
+                continue;
+            }
 
-        for (let i = 0; i < clusters.length; i++) {
-            if (i !== merge1 && i !== merge2) {
-                const clusterI = clusters[i];
-                if (!clusterI) {
-                    continue;
-                }
+            const dist = clusterDistance(newCluster, cluster, distances, linkage);
+            newDistMap.set(cluster.id, dist);
 
-                const newDist = clusterDistance(
-                    newCluster.members,
-                    clusterI.members,
-                    distances,
-                    linkage,
-                );
-                const newClusterDistances = clusterDistances.get(newCluster.id);
-                if (newClusterDistances) {
-                    newClusterDistances.set(clusterI.id, newDist);
-                }
+            // Update reverse direction
+            const clusterDistMap = clusterDistances.get(cluster.id);
+            if (clusterDistMap) {
+                clusterDistMap.set(newCluster.id, dist);
             }
         }
+        clusterDistances.set(newCluster.id, newDistMap);
 
-        // Remove old clusters and their distances
-        clusterDistances.delete(cluster1.id);
-        clusterDistances.delete(cluster2.id);
-        for (const [, dists] of clusterDistances) {
-            dists.delete(cluster1.id);
-            dists.delete(cluster2.id);
+        // Remove merged clusters
+        activeClusters.delete(mergeCluster1);
+        activeClusters.delete(mergeCluster2);
+        clusterDistances.delete(mergeCluster1.id);
+        clusterDistances.delete(mergeCluster2.id);
+
+        // Clean up references
+        for (const distMap of clusterDistances.values()) {
+            distMap.delete(mergeCluster1.id);
+            distMap.delete(mergeCluster2.id);
         }
 
-        // Remove merged clusters and add new one
-        clusters.splice(Math.max(merge1, merge2), 1);
-        clusters.splice(Math.min(merge1, merge2), 1);
-        clusters.push(newCluster);
+        activeClusters.add(newCluster);
     }
 
-    // Generate clusters at different levels
-    const clustersByLevel = new Map<number, Set<T>[]>();
+    // Root is the last remaining cluster
+    const root = activeClusters.size > 0 ? Array.from(activeClusters)[0] : dendrogram[dendrogram.length - 1];
 
-    function extractClustersAtHeight(node: ClusterNode<T>, height: number, clusters: Set<T>[]): void {
-        if (node.height <= height || !node.left || !node.right) {
-            clusters.push(node.members);
-        } else {
-            extractClustersAtHeight(node.left, height, clusters);
-            extractClustersAtHeight(node.right, height, clusters);
-        }
-    }
+    // Generate clusters at different heights
+    const clustersByLevel = new Map<number, Set<T>[]>();
 
-    // Handle the case where we have multiple disconnected components
-    let root: ClusterNode<T>;
-    if (clusters.length === 0) {
+    if (!root) {
         return {
             root: {id: "empty", members: new Set(), distance: 0, height: 0},
             dendrogram: [],
             clusters: new Map(),
         };
-    } else if (clusters.length === 1 && clusters[0]) {
-        root = clusters[0];
-    } else {
-        // Create a root that combines all remaining clusters (forest)
-        const allMembers = new Set<T>();
-        let maxHeight = 0;
-        for (const cluster of clusters) {
-            for (const member of cluster.members) {
-                allMembers.add(member);
-            }
-            maxHeight = Math.max(maxHeight, cluster.height);
-        }
-
-        root = {
-            id: "root-forest",
-            members: allMembers,
-            distance: Infinity,
-            height: maxHeight + 1,
-            // Store references to the individual trees
-            trees: [... clusters],
-        };
-
-        dendrogram.push(root);
     }
 
     for (let h = 0; h <= root.height; h++) {
-        const clustersAtHeight: Set<T>[] = [];
-
-        if (clusters.length > 1 && h === root.height) {
-            // At forest root height, return single cluster with all members
-            clustersAtHeight.push(root.members);
-        } else if (clusters.length > 1 && h < root.height) {
-            // For forest, extract from each tree separately
-            for (const tree of clusters) {
-                extractClustersAtHeight(tree, h, clustersAtHeight);
-            }
-        } else {
-            extractClustersAtHeight(root, h, clustersAtHeight);
-        }
-
-        clustersByLevel.set(h, clustersAtHeight);
+        clustersByLevel.set(h, cutDendrogram(root, h));
     }
 
     return {
@@ -332,7 +319,7 @@ export function hierarchicalClustering<T>(
 }
 
 /**
- * Cut dendrogram at specified height to get clusters
+ * Cut dendrogram at specific height to get clusters
  */
 export function cutDendrogram<T>(
     root: ClusterNode<T>,
@@ -341,27 +328,23 @@ export function cutDendrogram<T>(
     const clusters: Set<T>[] = [];
 
     function traverse(node: ClusterNode<T>): void {
-        // Handle forest root specially
-        if (node.trees) {
-            if (node.height <= height) {
-                // If we're at or below the requested height, return the whole forest as one cluster
-                clusters.push(node.members);
-            } else {
-                // Otherwise, traverse each tree in the forest
+        if (node.height <= height || (!node.left && !node.right)) {
+            clusters.push(node.members);
+        } else {
+            if (node.trees) {
+                // Handle forest nodes
                 for (const tree of node.trees) {
                     traverse(tree);
                 }
-            }
-
-            return;
-        }
+            } else {
+                if (node.left) {
+                    traverse(node.left);
+                }
 
-        // Normal binary tree traversal
-        if (node.height <= height || !node.left || !node.right) {
-            clusters.push(node.members);
-        } else {
-            traverse(node.left);
-            traverse(node.right);
+                if (node.right) {
+                    traverse(node.right);
+                }
+            }
         }
     }
 
@@ -370,7 +353,7 @@ export function cutDendrogram<T>(
 }
 
 /**
- * Cut dendrogram to get k clusters
+ * Cut dendrogram to get exactly k clusters
  */
 export function cutDendrogramKClusters<T>(
     root: ClusterNode<T>,
@@ -384,7 +367,7 @@ export function cutDendrogramKClusters<T>(
         return [root.members];
     }
 
-    // Find the height that gives us k clusters
+    // Binary search for the right height
     let low = 0;
     let high = root.height;
 
@@ -408,10 +391,11 @@ export function cutDendrogramKClusters<T>(
  * Compute modularity-based hierarchical clustering
  * Uses modularity gain to decide merges
  */
-export function modularityHierarchicalClustering<T>(
-    graph: Map<T, Set<T>>,
-): HierarchicalClusteringResult<T> {
-    const nodes = Array.from(graph.keys());
+export function modularityHierarchicalClustering(
+    graph: Graph,
+): HierarchicalClusteringResult<string> {
+    const adjacencySet = graphToAdjacencySet(graph);
+    const nodes = Array.from(adjacencySet.keys());
     const n = nodes.length;
 
     if (n === 0) {
@@ -422,117 +406,104 @@ export function modularityHierarchicalClustering<T>(
         };
     }
 
-    // Calculate total edges
-    let m = 0;
-    for (const neighbors of graph.values()) {
-        m += neighbors.size;
+    // Compute degrees
+    const degrees = new Map<string, number>();
+    let totalEdges = 0;
+
+    for (const [node, neighbors] of adjacencySet) {
+        degrees.set(node, neighbors.size);
+        totalEdges += neighbors.size;
     }
-    m = m / 2; // Each edge counted twice
+
+    // For undirected graphs
+    totalEdges = totalEdges / 2;
 
     // Initialize clusters
-    const clusters: ClusterNode<T>[] = nodes.map((node, i) => ({
+    const clusters: ClusterNode<string>[] = nodes.map((node, i) => ({
         id: `leaf-${String(i)}`,
         members: new Set([node]),
         distance: 0,
         height: 0,
     }));
 
-    const dendrogram: ClusterNode<T>[] = [... clusters];
+    const dendrogram: ClusterNode<string>[] = [... clusters];
     let clusterCount = n;
 
-    // Track community assignments
-    const communityMap = new Map<T, number>();
-    nodes.forEach((node, i) => communityMap.set(node, i));
-
-    // Compute modularity gain for merging communities
-    function modularityGain(comm1: Set<T>, comm2: Set<T>): number {
-        let edgesBetween = 0;
-        let degree1 = 0;
-        let degree2 = 0;
+    // Active clusters
+    const activeClusters = new Set(clusters);
 
-        for (const node1 of comm1) {
-            const neighbors = graph.get(node1);
-            if (!neighbors) {
-                continue;
-            }
+    // Community assignments
+    const communities = new Map<string, number>();
+    nodes.forEach((node, i) => communities.set(node, i));
 
-            degree1 += neighbors.size;
+    // Merge clusters based on modularity gain
+    while (activeClusters.size > 1) {
+        let maxGain = -Infinity;
+        let mergeCluster1: ClusterNode<string> | null = null;
+        let mergeCluster2: ClusterNode<string> | null = null;
 
-            for (const neighbor of neighbors) {
-                if (comm2.has(neighbor)) {
-                    edgesBetween++;
+        for (const cluster1 of activeClusters) {
+            for (const cluster2 of activeClusters) {
+                if (cluster1.id >= cluster2.id) {
+                    continue;
                 }
-            }
-        }
 
-        for (const node2 of comm2) {
-            const node2Neighbors = graph.get(node2);
-            if (node2Neighbors) {
-                degree2 += node2Neighbors.size;
-            }
-        }
+                // Calculate modularity gain
+                let edgesBetween = 0;
+                let degreeProduct = 0;
 
-        // Modularity gain formula
-        const gain = (edgesBetween / m) - ((degree1 * degree2) / (4 * m * m));
-        return gain;
-    }
+                for (const node1 of cluster1.members) {
+                    const neighbors = adjacencySet.get(node1);
+                    const degree1 = degrees.get(node1) ?? 0;
 
-    // Merge clusters based on modularity
-    while (clusters.length > 1) {
-        let maxGain = -Infinity;
-        let merge1 = 0;
-        let merge2 = 0;
-
-        for (let i = 0; i < clusters.length; i++) {
-            for (let j = i + 1; j < clusters.length; j++) {
-                const clusterI = clusters[i];
-                const clusterJ = clusters[j];
-                if (!clusterI || !clusterJ) {
-                    continue;
+                    for (const node2 of cluster2.members) {
+                        if (neighbors?.has(node2)) {
+                            edgesBetween++;
+                        }
+
+                        const degree2 = degrees.get(node2) ?? 0;
+                        degreeProduct += degree1 * degree2;
+                    }
                 }
 
-                const gain = modularityGain(clusterI.members, clusterJ.members);
+                const gain = (edgesBetween / totalEdges) -
+                    (degreeProduct / (4 * totalEdges * totalEdges));
 
                 if (gain > maxGain) {
                     maxGain = gain;
-                    merge1 = i;
-                    merge2 = j;
+                    mergeCluster1 = cluster1;
+                    mergeCluster2 = cluster2;
                 }
             }
         }
 
-        // Create new cluster
-        const cluster1 = clusters[merge1];
-        const cluster2 = clusters[merge2];
-        if (!cluster1 || !cluster2) {
-            continue;
+        if (!mergeCluster1 || !mergeCluster2) {
+            break;
         }
 
-        const newCluster: ClusterNode<T> = {
+        // Create new cluster
+        const newCluster: ClusterNode<string> = {
             id: `cluster-${String(clusterCount++)}`,
-            members: new Set([... cluster1.members, ... cluster2.members]),
-            left: cluster1,
-            right: cluster2,
+            members: new Set([... mergeCluster1.members, ... mergeCluster2.members]),
+            left: mergeCluster1,
+            right: mergeCluster2,
             distance: -maxGain, // Use negative gain as distance
-            height: Math.max(cluster1.height, cluster2.height) + 1,
+            height: Math.max(mergeCluster1.height, mergeCluster2.height) + 1,
         };
 
         dendrogram.push(newCluster);
 
-        // Update community assignments
-        for (const node of newCluster.members) {
-            communityMap.set(node, clusterCount - 1);
-        }
-
-        // Remove merged clusters
-        clusters.splice(Math.max(merge1, merge2), 1);
-        clusters.splice(Math.min(merge1, merge2), 1);
-        clusters.push(newCluster);
+        // Remove old and add new
+        activeClusters.delete(mergeCluster1);
+        activeClusters.delete(mergeCluster2);
+        activeClusters.add(newCluster);
     }
 
-    // Generate clusters at different levels
-    const clustersByLevel = new Map<number, Set<T>[]>();
-    const root = clusters[0];
+    // Root is the last remaining cluster
+    const root = Array.from(activeClusters)[0] ?? dendrogram[dendrogram.length - 1];
+
+    // Generate clusters at different heights
+    const clustersByLevel = new Map<number, Set<string>[]>();
 
     if (!root) {
         return {
@@ -553,3 +524,22 @@ export function modularityHierarchicalClustering<T>(
     };
 }
 
+/**
+ * Agglomerative hierarchical clustering
+ * Builds clusters bottom-up by merging closest pairs
+ *
+ * @param graph - Undirected graph - accepts Graph class or Map<T, Set<T>>
+ * @param linkage - Linkage method for cluster distance
+ * @returns Hierarchical clustering result
+ *
+ * Time Complexity: O(n³) naive, O(n² log n) with heap
+ * Space Complexity: O(n²)
+ */
+export function hierarchicalClustering(
+    graph: Graph,
+    linkage: LinkageMethod = "single",
+): HierarchicalClusteringResult<string> {
+    // Convert Graph to adjacency set representation
+    const adjacencySet = graphToAdjacencySet(graph);
+    return hierarchicalClusteringImpl(adjacencySet, linkage);
+}