Graph.scala

Graph.scala

// https://github.com/asarkar/algorithms-design-analysis-2/blob/master/data-structures/src/main/scala/org/asarkar/data/Graph.scala

import scala.collection.immutable.{HashMap, Set}

sealed trait Edge[T] {
  def tail: T
  def head: T
}

sealed trait WeightedEdge[T] extends Edge[T] {
  def weight: Double
}

sealed abstract class AbstractUndirectedEdge[T] extends Edge[T] {
  def other(v: T): Option[T] = if (v == tail) Some(head) else if (v == head) Some(tail) else None

  override def equals(obj: Any): Boolean = {
    obj match {
      case that: UndirectedEdge[T] => this.hashCode == that.hashCode
      case _ => false
    }
  }

  override def hashCode(): Int = {
    Seq(head, tail)
      .map(_.hashCode)
      .sorted
      .foldLeft(7)((result, hc) => 31 * result + hc)
  }
}

case class UndirectedEdge[T](override val tail: T, override val head: T)
  extends AbstractUndirectedEdge[T]

case class UndirectedWeightedEdge[T](override val tail: T, override val head: T, override val weight: Double)
  extends AbstractUndirectedEdge[T] with WeightedEdge[T]

case class DirectedEdge[T](override val tail: T, override val head: T)
  extends Edge[T]

case class DirectedWeightedEdge[T](override val tail: T, override val head: T, override val weight: Double)
  extends WeightedEdge[T]

sealed trait Graph[V, E <: Edge[V]] {
  def addEdge(edge: E): Graph[V, E]

  def vertices: Iterable[V]

  def outgoing(v: V): Iterable[E]

  def contains(v: V): Boolean

  def edges: Iterable[E] = vertices.flatMap(outgoing).toSet
  def hasEdge(v: V, w: V): Boolean = neighbors(v).exists(_ == w)
  def neighbors(v: V): Iterable[V] = outgoing(v).map(e => if (e.head == v) e.tail else e.head)

  def dump(): Unit
}

sealed private abstract class BaseGraph[V, E <: Edge[V]](protected val _vertices: HashMap[V, Set[E]]) extends Graph[V, E] {
  protected def addEdge0(_vs: HashMap[V, Set[E]], edge: E): HashMap[V, Set[E]] = {
    if (_vs.exists(entry => entry._1 == edge.tail)) {
      _vs.map(e => if (e._1 != edge.tail) e else (edge.tail, e._2 + edge))
    } else {
      _vs + (edge.tail -> Set(edge))
    }
  }


  override def vertices: Iterable[V] = _vertices.keys
  override def outgoing(v: V): Iterable[E] = Option.option2Iterable(_vertices.get(v)).flatten
  override def contains(v: V): Boolean = _vertices.contains(v)
}

sealed private class UndirectedGraph[V, E <: Edge[V]](private val _vs: HashMap[V, Set[E]] = new HashMap[V, Set[E]]()) extends BaseGraph[V, E](_vs) {

  override def addEdge(edge: E): Graph[V, E] = {
    val nextVs = addEdge0(_vs, edge)

    val reverse = edge match {
      case e: UndirectedWeightedEdge[V] => e.copy(tail = e.head, head = e.tail)
      case e => UndirectedEdge(e.head, e.tail)
    }

    val nextVsReverse = addEdge0(nextVs, reverse.asInstanceOf[E])
    new UndirectedGraph(nextVsReverse)
  }

  def dump(): Unit = println(_vs)
}

sealed private class DirectedGraph[V, E <: Edge[V]](private val _vs: HashMap[V, Set[E]] = new HashMap[V, Set[E]]()) extends BaseGraph[V, E](_vs) {

  override def addEdge(edge: E): Graph[V, E] = {
    val nextVs = addEdge0(_vs, edge)
    new DirectedGraph(nextVs)
  }

  def dump(): Unit = println(_vs)
}

object Graph {
  def undirectedBuilder[V, E <: AbstractUndirectedEdge[V]](): Graph[V, E] = new UndirectedGraph[V, E]()

  def directedBuilder[V, E <: DirectedEdge[V]](): Graph[V, E] = new DirectedGraph[V, E]()
}