Edge Weighted

import scala.collection.mutable.{LinkedHashSet, SortedMap, PriorityQueue}
import scala.annotation.tailrec

sealed abstract class EdgeWeightedGraph[A](m: Map[A, Set[(A, Double)]]) {

    object weighteduf {

      type Id = Map[A, A]
      type Sz = Map[A, Int]

      val (id, sz) = m.keys.foldLeft(Map.empty[A, A] -> Map.empty[A, Int]) {
        case ((id, sz), k) => (id + (k -> k)) -> (sz + (k -> 1))
      }

      @tailrec
      def find(a: A, m: Id): A = {
        if (m(a) == a) a
        else find(m(a), m)
      }

      def connected(a: A, b: A, m: Id): Boolean =
        find(a, m) == find(b, m)

      def union(
          a: A,
          b: A,
          id: Id,
          sz: Sz
      ): (Id, Sz) = {

        val fa = find(a, id)
        val fb = find(b, id)
        if (fa == fb) {
          id -> sz
        } else {
          if (sz(fa) < sz(fb))
            id.updated(fa, fb) -> sz.updated(fb, sz(fb) + sz(fa))
          else id.updated(fb, fa) -> sz.updated(fa, sz(fa) + sz(fb))
        }

      }

    }

    def prims(start: A): Set[(A, A)] = {

      val keys = m.keySet

      implicit val ordering: Ordering[(A, A, Double)] =
        new Ordering[(A, A, Double)] {
          override def compare(x: (A, A, Double), y: (A, A, Double)): Int =
            implicitly[Ordering[Double]].compare(x._3, y._3) * -1
        }

      def go(
          queue: PriorityQueue[(A, A, Double)],
          added: Set[(A, A)],
          visited: Set[A],
          invalidated: Set[Set[A]]
      ): (Set[(A, A)], Set[A]) =
        if (visited.size == keys.size)
          added -> visited
        else {

          val queue1 = queue
            .dropWhile {
              case (v1, v2, _) => invalidated.contains(Set(v1, v2))
            }

          queue1.headOption
            .fold(
              added -> visited
            ) {
              case (v1, v2, _) =>
                val (nexts, inval) = m
                  .get(v2)
                  .fold(
                    (List.empty[(A, A, Double)], Set.empty[Set[A]])
                  ) { nextEdges =>
                    val (inval, next) = nextEdges.partition {
                      case (v, _) => visited.contains(v)
                    }
                    (next.map { case (v, distance) => (v2, v, distance) }.toList,
                     inval.map {
                       case (v, _) => Set(v2, v)
                     })
                  }

                queue1.dequeue()
                queue1.enqueue(nexts: _*)

                go(
                  queue1,
                  added + (v1 -> v2),
                  visited + v2,
                  invalidated ++ inval
                )
            }
        }

      val pq = new PriorityQueue[(A, A, Double)]()
      pq.enqueue(
        m.get(start)
          .fold(List.empty[(A, A, Double)])(_.map {
            case (v, distance) => (start, v, distance)
          }.toList): _*
      )

      go(pq, Set.empty, Set(start), Set.empty)._1
    }

    def kruskal(): Set[(A, A)] = {

      implicit val ord: Ordering[(Double, A, A)] =
        new Ordering[(Double, A, A)] {
          override def compare(x: (Double, A, A), y: (Double, A, A)): Int =
            implicitly[Ordering[Double]].compare(x._1, y._1) * -1
        }

      val (_, pq) =
        m.foldLeft(Set.empty[Set[A]] -> PriorityQueue.empty[(Double, A, A)]) {
          case ((added, pq), (edge1, edgeWeightSet)) =>
            edgeWeightSet.foldLeft(added -> pq) {
              case ((added, sm), (edge2, weight)) =>
                if (added.contains(Set(edge1, edge2))) added -> sm
                else
                  added + Set(edge1, edge2) ->
                    (pq += ((weight, edge1, edge2)))
            }
        }

      @tailrec
      def go(
          pq: PriorityQueue[(Double, A, A)],
          id: weighteduf.Id,
          sz: weighteduf.Sz,
          mst: Set[(A, A)]
      ): Set[(A, A)] =
        if (pq.isEmpty) mst
        else {
          val (_, e1, e2) = pq.dequeue()
          if (weighteduf.connected(e1, e2, id)) go(pq, id, sz, mst)
          else {
            val (id1, sz1) = weighteduf.union(e1, e2, id, sz)
            go(pq, id1, sz1, (mst + (e1 -> e2)))
          }
        }

      go(pq, weighteduf.id, weighteduf.sz, Set.empty[(A, A)])

    }

}