コード
let input = require('fs').readFileSync('/dev/stdin', 'utf8');
const lines = input.trim().split("\n")
let v, e, path = []
for (let i = 0; i < lines.length; i++) {
if (i === 0) {
[v, e] = lines[i].split(" ").map(v => parseInt(v))
path = [...Array(v)].map(() => [])
continue
}
const [s, t, w] = lines[i].split(" ").map(v => parseInt(v))
path[s].push([t, w])
path[t].push([s, w])
}
const parent = (i) => Math.floor((i - 1) / 2)
const leftChild = (i) => 2 * i + 1
const rightChild = (i) => 2 * i + 2
const MinHeap = function (n) {
this.a = [...Array(n).keys()].map(v => [v, Number.MAX_SAFE_INTEGER]) // [vertex, weight]
this.pos = [...Array(n).keys()].map(v => v)
this.size = n
}
MinHeap.prototype.extractMin = function () {
const [v, w] = this.a[0]
this.swap(0, this.size - 1)
this.size--
this.popDown(0)
return [v, w]
}
MinHeap.prototype.swap = function (i, j) {
const vi = this.a[i][0], vj = this.a[j][0];
[this.a[i], this.a[j]] = [this.a[j], this.a[i]];
[this.pos[vi], this.pos[vj]] = [this.pos[vj], this.pos[vi]]
}
MinHeap.prototype.decreaseKey = function (v, val) {
const i = this.pos[v]
this.a[i][1] = val
this.popUp(i)
}
// vertexが含まれているか
MinHeap.prototype.contains = function (v) {
return this.pos[v] < this.size
}
// 適切な位置まで上げる
MinHeap.prototype.popUp = function (i) {
let pi = parent(i)
while (i > 0 && this.a[i][1] < this.a[pi][1]) { // 親が子より大きかったら交換
this.swap(i, pi)
i = pi
pi = parent(i)
}
}
// 値を取得
MinHeap.prototype.val = function (v) {
return this.a[this.pos[v]][1]
}
// 適切な位置まで下げる
MinHeap.prototype.popDown = function (i) {
const val = this.a[i][1]
let candidate = null // 子供のほうが小さかったら交換
const ri = rightChild(i)
const li = leftChild(i)
if (ri < this.size) { // right child exist
if (this.a[ri][1] < val) {
candidate = ri
}
}
if (li < this.size) {
if (candidate == null) {
if (this.a[li][1] < val) {
candidate = li
}
} else {
if (this.a[li][1] < this.a[ri][1]) {
candidate = li
}
}
}
if (candidate != null) {
this.swap(i, candidate)
this.popDown(candidate)
}
}
function prim(n, path) {
const minHeap = new MinHeap(n)
minHeap.decreaseKey(0, 0)
let total = 0
while (minHeap.size > 0) {
const [v, w] = minHeap.extractMin()
// check neighbors
for (const [next, nextWeight] of path[v]) {
if (!minHeap.contains(next)) {
continue
}
const currentWeight = minHeap.val(next)
if (nextWeight < currentWeight) {
minHeap.decreaseKey(next, nextWeight)
}
}
total += w
}
console.log(total)
}
// console.log(v, e)
// console.log(path)
prim(v, path)
Top comments (0)