TypeScript implementation of disjoint set (union and find algorithm)

disjointSet.ts

export function createDisjointSet(initialSize: number) {
    const roots = Array.from(new Array(initialSize), (_, index) => index)
    const sizes = roots.map(() => 1)
    let groups = initialSize

    // return the id of the set which given index-th element exists.
    // but set id may changes if further union operations are done.
    function find(index: number): number {
        if (roots[index] === index) return index
        return roots[index] = find(roots[index])
    }

    function union(ia: number, ib: number): void {
        let ra = find(ia)
        let rb = find(ib)

        if (ra === rb) return

        if (sizes[ra] < sizes[rb]) {
            [ra, rb] = [rb, ra]
        }
        
        roots[rb] = ra
        sizes[ra] += sizes[rb]
        groups -= 1
    }

    // return the size of the set which has index-th element.
    // if index is not given, return the number of dijoint sets.
    function size(index?: number): number {
        if (index === undefined) return groups
        return sizes[find(index)]
    }

    return {
        find,
        union,
        size,
    }
}

Simple disjoint set implementation which can only obtain the root of the sets and size of them. This uses path halving and merge by size strategies suggested by Tarjan and Jan van Leeuwen, to obtain O(α(n)) time complexity for every operation asymptotically. (Here α is Inverse Ackermann Function which insanely slowly increases)

For more information, checkout the wiki