UnionFind2.js

function UnionFind(n) {
    this.id = [...Array(n)].map((_, i) => i)
    this.sz = [...Array(n)].map(_ => 1)
    this.max = [...this.id]
}

UnionFind.prototype._root = function (id) {
    let max = id
    while(id != this.id[id]) {
        id = this.id[id]
        max = Math.max(max, id)
    }
    return [id, max]
}

UnionFind.prototype.union = function (i, k) {
    let [iroot, imax] = this._root(i)
    let [kroot, kmax] = this._root(k)

    if (iroot === kroot) return

    if (this.sz[iroot] > this.sz[kroot]) {
        this.id[kroot] = iroot
        this.sz[iroot] += this.sz[kroot]
    } else {
        this.id[iroot] = kroot
        this.sz[kroot] += this.sz[iroot]
    }

    this.max[iroot] = this.max[kroot] = Math.max(imax, kmax)
}

UnionFind.prototype.findMax = function (i) {
    return this.max[this._root(i)[0]]
}