graph.clj

(ns ^{:doc "Defines protocols for graphs, digraphs, and weighted graphs.

Also provides record implementations and constructors for simple graphs --
weighted, unweighted, directed, and undirected. The implementations are based
on adjacency lists."
      :author "Justin Kramer"}
  loom.graph)

(set! *warn-on-reflection* true)

;;;
;;; Protocols
;;;

(defprotocol Graph
  (add-nodes* [g nodes] "Add nodes to graph g. Non-variadic; see add-nodes")
  (add-edges* [g edges] "Add edges to graph g. Non-variadic; See add-edges")
  (remove-nodes* [g nodes] "Remove nodes from graph g. Non-variadic; see remove-nodes")
  (remove-edges* [g nodes] "Removes edges from graph g. Non-variadic; see remove-edges")
  (remove-all [g] "Removes all nodes and edges from graph g")
  (nodes [g] "Return a collection of the nodes in graph g")
  (edges [g] "Return a collection of the edges in graph g")
  (has-node? [g node] "Return true when node is in g")
  (has-edge? [g n1 n2] "Return true when edge [n1 n2] is in g")
  (neighbors [g] [g node] "Return nodes adjacent to node, or (partial neighbors g)"))

(defprotocol Digraph
  (incoming [g node] "Return nodes pointing towards node")
  (transpose [g] "Return a graph with all edges reversed"))

(defprotocol WeightedGraph
  (weight [g] [g n1 n2] "Return weight of edge [n1 n2], or (partial weight g)"))

;;;
;;; Records
;;;

(defrecord SimpleGraph [nodeset adj])
(defrecord SimpleDigraph [nodeset adj in])
(defrecord SimpleWeightedGraph [nodeset adj])
(defrecord SimpleWeightedDigraph [nodeset adj in])

(def ^{:doc "Weight used when none is given for edges in weighted graphs"}
  *default-weight* 1)

(def default-graph-impls
  {:all
   {:nodes (fn [g]
             (:nodeset g))
    :edges (fn [g]
             (for [n1 (nodes g)
                   n2 (neighbors g n1)]
               [n1 n2]))
    :has-node? (fn [g node]
                 (contains? (:nodeset g) node))
    :has-edge? (fn [g n1 n2]
                 (contains? (get-in g [:adj n1]) n2))}

   ;; Unweighted graphs store adjacencies as {node #{neighbor}}
   :unweighted
   {:add-nodes* (fn [g nodes]
                  (reduce
                   (fn [g n]
                     (-> g
                         (update-in [:nodeset] conj n)
                         (assoc-in [:adj n] (or ((:adj g) n) #{}))))
                   g nodes))
    :neighbors (fn
                 ([g] (partial neighbors g))
                 ([g node] (get-in g [:adj node])))}
   
   ;; Weighted graphs store adjacencies as {node {neighbor weight}}
   :weighted
   {:add-nodes* (fn [g nodes]
                  (reduce
                   (fn [g n]
                     (-> g
                         (update-in [:nodeset] conj n)
                         (assoc-in [:adj n] (or ((:adj g) n) {}))))
                   g nodes))
    :neighbors (fn
                 ([g] (partial neighbors g))
                 ([g node] (keys (get-in g [:adj node]))))}})

(def default-digraph-impl
  {:incoming (fn [g node]
               (get-in g [:in node]))})

(def default-weighted-graph-impl
  {:weight (fn
             ([g] (partial weight g))
             ([g n1 n2] (get-in g [:adj n1 n2])))})

(defn- remove-adj-nodes [m nodes adjacents remove-fn]
  (reduce
   (fn [m n]
     (if (m n)
       (update-in m [n] #(apply remove-fn % nodes))
       m))
   (apply dissoc m nodes)
   adjacents))

(extend SimpleGraph
  Graph
  (assoc (apply merge (map default-graph-impls [:all :unweighted]))

    :add-edges*
    (fn [g edges]
      (reduce
       (fn [g [n1 n2]]
         (-> g
             (update-in [:nodeset] conj n1 n2)
             (update-in [:adj n1] (fnil conj #{}) n2)
             (update-in [:adj n2] (fnil conj #{}) n1)))
       g edges))
    
    :remove-nodes*
    (fn [g nodes]
      (let [nbrs (mapcat #(neighbors g %) nodes)]
        (-> g
            (update-in [:nodeset] #(apply disj % nodes))
            (assoc :adj (remove-adj-nodes (:adj g) nodes nbrs disj)))))
    
    :remove-edges*
    (fn [g edges]
      (reduce
       (fn [g [n1 n2]]
         (-> g
             (update-in [:adj n1] disj n2)
             (update-in [:adj n2] disj n1)))
       g edges))

    :remove-all
    (fn [g]
      (assoc g :nodeset #{} :adj {}))))

(extend SimpleDigraph
  Graph
  (assoc (apply merge (map default-graph-impls [:all :unweighted]))
    
    :add-edges*
    (fn [g edges]
      (reduce
       (fn [g [n1 n2]]
         (-> g
             (update-in [:nodeset] conj n1 n2)
             (update-in [:adj n1] (fnil conj #{}) n2)
             (update-in [:in n2] (fnil conj #{}) n1)))
       g edges))
    
    :remove-nodes*
    (fn [g nodes]
      (let [ins (mapcat #(incoming g %) nodes)
            outs (mapcat #(neighbors g %) nodes)]
        (-> g
            (update-in [:nodeset] #(apply disj % nodes))
            (assoc :adj (remove-adj-nodes (:adj g) nodes ins disj))
            (assoc :in (remove-adj-nodes (:in g) nodes outs disj)))))
    
    :remove-edges*
    (fn [g edges]
      (reduce
       (fn [g [n1 n2]]
         (-> g
             (update-in [:adj n1] disj n2)
             (update-in [:in n2] disj n1)))
       g edges))

    :remove-all
    (fn [g]
      (assoc g :nodeset #{} :adj {} :in {})))

  Digraph
  (assoc default-digraph-impl
    :transpose (fn [g]
                 (assoc g :adj (:in g) :in (:adj g)))))

(extend SimpleWeightedGraph
  Graph
  (assoc (apply merge (map default-graph-impls [:all :weighted]))

    :add-edges*
    (fn [g edges]
      (reduce
       (fn [g [n1 n2 & [w]]]
         (-> g
             (update-in [:nodeset] conj n1 n2)
             (assoc-in [:adj n1 n2] (or w *default-weight*))
             (assoc-in [:adj n2 n1] (or w *default-weight*))))
       g edges))
    
    :remove-nodes*
    (fn [g nodes]
      (let [nbrs (mapcat #(neighbors g %) nodes)]
        (-> g
            (update-in [:nodeset] #(apply disj % nodes))
            (assoc :adj (remove-adj-nodes (:adj g) nodes nbrs dissoc)))))
    
    :remove-edges*
    (fn [g edges]
      (reduce
       (fn [g [n1 n2]]
         (-> g
             (update-in [:adj n1] dissoc n2)
             (update-in [:adj n2] dissoc n1)))
       g edges))

    :remove-all
    (fn [g]
      (assoc g :nodeset #{} :adj {})))
  
  WeightedGraph
  default-weighted-graph-impl)

(extend SimpleWeightedDigraph
  Graph
  (assoc (apply merge (map default-graph-impls [:all :weighted]))
    
    :add-edges*
    (fn [g edges]
      (reduce
       (fn [g [n1 n2 & [w]]]
         (-> g
             (update-in [:nodeset] conj n1 n2)
             (assoc-in [:adj n1 n2] (or w *default-weight*))
             (update-in [:in n2] (fnil conj #{}) n1)))
       g edges))
    
    :remove-nodes*
    (fn [g nodes]
      (let [ins (mapcat #(incoming g %) nodes)
            outs (mapcat #(neighbors g %) nodes)]
        (-> g
            (update-in [:nodeset] #(apply disj % nodes))
            (assoc :adj (remove-adj-nodes (:adj g) nodes ins dissoc))
            (assoc :in (remove-adj-nodes (:in g) nodes outs disj)))))
    
    :remove-edges*
    (fn [g edges]
      (reduce
       (fn [g [n1 n2]]
         (-> g
             (update-in [:adj n1] dissoc n2)
             (update-in [:in n2] dissoc n1)))
       g edges))

    :remove-all
    (fn [g]
      (assoc g :nodeset #{} :adj {}) :in {}))

  Digraph
  (assoc default-digraph-impl
    :transpose (fn [g]
                 (reduce (fn [tg [n1 n2]]
                             (add-edges* tg [[n2 n1 (weight g n1 n2)]]))
                         (assoc g :adj {} :in {})
                         (edges g))))
  
  WeightedGraph
  default-weighted-graph-impl)

;;;
;;; Utility functions and constructors
;;;

(defn add-nodes
  "Add nodes to graph g. Nodes can be any type of object"
  [g & nodes]
  (add-nodes* g nodes))

(defn add-edges
  "Add edges to graph g. For unweighted graphs, edges take the form [n1 n2].
  For weighted graphs, edges take the form [n1 n2 weight] or [n1 n2], the
  latter defaulting to a weight of 1"
  [g & edges]
  (add-edges* g edges))

(defn remove-nodes
  "Remove nodes from graph g"
  [g & nodes]
  (remove-nodes* g nodes))

(defn remove-edges
  "Remove edges from graph g. Do not include weights"
  [g & edges]
  (remove-edges* g edges))

(defn build-graph
  "Builds up a graph (i.e. adds edges and nodes) from any combination of
  other graphs, adjacency maps, edges, and nodes."
  [g & inits]
  (letfn [(build [g init]
            (cond
             ;; graph
             (satisfies? Graph init)
             (if (and (satisfies? WeightedGraph g)
                      (satisfies? WeightedGraph init))
               (reduce add-edges
                       (add-nodes* g (nodes init))
                       (for [[n1 n2] (edges init)]
                         [n1 n2 (weight init n1 n2)]))
               (-> g
                   (add-nodes* (nodes init))
                   (add-edges* (edges init))))
             ;; adacency map
             (map? init)
             (let [es (if (map? (val (first init)))
                        (for [[n nbrs] init
                              [nbr wt] nbrs]
                          [n nbr wt])
                        (for [[n nbrs] init
                              nbr nbrs]
                          [n nbr]))]
               (reduce add-edges g es))
             ;; edge
             (sequential? init) (add-edges g init)
             ;; node
             :else (add-nodes g init)))]
    (reduce build g inits)))

(defn graph
  "Create an unweighted, undirected graph. inits can be edges, adjacency maps,
  or graphs"
  [& inits]
  (apply build-graph (SimpleGraph. #{} {}) inits))

(defn digraph
  "Create an unweighted, directed graph. inits can be edges, adjacency maps,
  or graphs"
  [& inits]
  (apply build-graph (SimpleDigraph. #{} {} {}) inits))

(defn weighted-graph
  [& inits]
  "Create an weighted, undirected graph. inits can be edges, adjacency maps,
  or graphs"
  (apply build-graph (SimpleWeightedGraph. #{} {}) inits))

(defn weighted-digraph
  "Create an weighted, directed graph. inits can be edges, adjacency maps,
  or graphs"
  [& inits]
  (apply build-graph (SimpleWeightedDigraph. #{} {} {}) inits))