Sunday, January 6, 2013

Strongly connected components in Clojure

Recently I had fun implementing strongly connected components algorithm in Clojure. I thought some of the Clojurians might want to see it.

(defn dfs
  "Depth first search. Short form of the method passes through all the 
  nodes of the graph even if it's disconnected .
  (nodes-fn graph) expected to return list of all the nodes in the graph.
  (child-fn graph node) expected to return list of all the nodes linked
   to the given node.
  Returns hash-map where nodes are associated with a pair :idx, :leader.
  :idx stores finishing index of the node traversal (post-order counter)
  :leader first finishing index of the current DFS."
  ([graph nodes-fn child-fn]
     (second
      (reduce ;; Start DFS from each node of the graph
       (fn [[idx result passed :as args] next-node]
         (if (not (passed next-node)) ;; Don't do DFS if node is marked
           (dfs idx idx result passed graph next-node child-fn)
           args))
       [0 {} #{}] ;;Initial index, result, set of passed nodes
       (nodes-fn graph))))
  ([idx leader result passed graph node child-fn]
     (let [[idx result passed]
           (reduce (fn [[idx result passed :as args] child-node]
                     (if (not (passed child-node))
                       (dfs idx leader result passed 
                            graph child-node child-fn)
                       args))
                   [idx result (conj passed node)]
                   (child-fn graph node))]
       [(inc idx)
        (assoc result node {:idx idx :leader leader})
        passed])))

(defn pass-two 
  "Calls DFS making sure that traversal is done in the reverse :idx order."
  [graph result child-fn]
  (let [nodes-fn 
        (constantly (->> result 
                         ;;Sort by :idx in reverse order
                         (sort-by (comp :idx second)) reverse 
                         ;;Return only nodes
                         (map first)))]
    (dfs graph nodes-fn child-fn)))
  
(defn scc 
  "Finds strongly connected components of the given directed graph.
  Returns lists of nodes grouped into SCC.
  (nodes-fn graph) expected to return list of all the nodes in the graph.
  (incoming-fn graph node) expected to return all the nodes with
   transitions towards the given node.
  (outgoing-fn graph node) expected to return all the nodes with
   transitions from the given node."
  [graph nodes-fn incoming-fn outgoing-fn]
  (let [result (dfs graph nodes-fn incoming-fn)
        leaders-idx (pass-two graph result outgoing-fn)]
    (for [scc-group (vals (group-by (comp :leader second) leaders-idx))]
      (for [[node & _] scc-group] node))))

Implementation is quite generic and can be used for different graph implementations. Here's an example using list multimap. I use the same graph as in the course given by Tim Roughgarden on Coursera:

(defn list-multimap 
  "Builds list multimap: {key1 [val1 val2 ...], key2 [val3 val4 ...]}.
   Each call adds value to the list associated with the key."
  [m [k v]]
  (if (m k)
    (update-in m [k] conj v)
    (assoc m k [v])))

(defn reverse-graph 
  "Reverses list multimap based graph, see below."
  [graph]
  (reduce
   list-multimap
   {}
   (for [[key values] graph v values] [v key])))

(def test-graph 
  {6 [9], 2 [8], 4 [7], 3 [6], 8 [5 6], 1 [4], 9 [3 7], 5 [2], 7 [1]})

;Same graph but with letters in case you confuse indices and node names.
;(def test-graph
;  {'a ['g], 'b ['e], 'c ['i], 'd ['a], 
;   'e ['h], 'f ['c 'h], 'g ['d 'i], 'h ['b], 'i ['f]})

(def reverse-test-graph
  (reverse-graph test-graph))

And here's how to use the dfs and scc functions:

(dfs test-graph 
     ;;fn that returns set of nodes
     (constantly (into #{} (flatten (seq test-graph))))
     ;;(get graph node) returns list of related nodes.
     get) 

;{2 {:idx 8, :leader 3},
; 8 {:idx 7, :leader 3},
; 6 {:idx 6, :leader 3},
; 9 {:idx 5, :leader 3},
; 3 {:idx 4, :leader 3},
; 5 {:idx 3, :leader 3},
; 1 {:idx 2, :leader 0},
; 4 {:idx 1, :leader 0},
; 7 {:idx 0, :leader 0}}

(scc test-graph 
     ;;fn that returns set of nodes
     (constantly (into #{} (flatten (seq test-graph))))
     ;;works as incoming-fn using cashed reversed graph
     #(get reverse-test-graph %2) 
     ;;(get graph node) returns list of related nodes
     get)

;((8 5 2) (9 3 6) (1 4 7))

Full source code is here.