prim.js 1.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. const _ = require('../lodash')
  2. const Graph = require('../graph')
  3. const PriorityQueue = require('../data/priority-queue')
  4. module.exports = prim
  5. function prim (g, weightFunc) {
  6. const result = new Graph()
  7. const parents = {}
  8. const pq = new PriorityQueue()
  9. let v
  10. function updateNeighbors (edge) {
  11. const w = edge.v === v ? edge.w : edge.v
  12. const pri = pq.priority(w)
  13. if (pri !== undefined) {
  14. var edgeWeight = weightFunc(edge)
  15. if (edgeWeight < pri) {
  16. parents[w] = v
  17. pq.decrease(w, edgeWeight)
  18. }
  19. }
  20. }
  21. if (g.nodeCount() === 0) {
  22. return result
  23. }
  24. _.each(g.nodes(), function (v) {
  25. pq.add(v, Number.POSITIVE_INFINITY)
  26. result.setNode(v)
  27. })
  28. // Start from an arbitrary node
  29. pq.decrease(g.nodes()[0], 0)
  30. var init = false
  31. while (pq.size() > 0) {
  32. v = pq.removeMin()
  33. if (_.has(parents, v)) {
  34. result.setEdge(v, parents[v])
  35. } else if (init) {
  36. throw new Error('Input graph is not connected: ' + g)
  37. } else {
  38. init = true
  39. }
  40. g.nodeEdges(v).forEach(updateNeighbors)
  41. }
  42. return result
  43. }