LEARNING JOURNAL

/**
 * Definition for a category handler.
 * class CategoryHandler {
 *     constructor(categories: number[]);
 *     public haveSameCategory(a: number, b: number): boolean;
 * }
 */
function numberOfCategories(n: number, categoryHandler: CategoryHandler): number {
    // To find number of unique categories, we need union-find algorithm.
    const uF = new UnionFind(n);

    // Go through each pair of elements
    for(let i = 0; i < n; i++) {
        for(let j = i + 1; j < n; j++) {
            // For each pair, check if they belong to the same category
            if(categoryHandler.haveSameCategory(i,j)){
                // Union them
                uF.union(i,j);
            }
        }
    }

    // return the number of categories
    return uF.getNumCategories();
};

// Union-Find algorithm implementation
// We need something to keep track of the categories' roots
// --> best way to do it is to create a class that has 
//     property that can be used to keep track of the roots
// In JS or TS, you can use function, but class looks cleanner.

class UnionFind {
    private roots: number[];
    private catSizes: number[];
    private numCategories:number;

    constructor(n:number) {
        this.roots = Array.from({length: n}).map((_,i) => i);
        this.catSizes = Array.from({length:n}).map(() => 1);
        this.numCategories = n;
    }

    public union(a:number, b:number) {
        // make them being at the same category
        const rootA = this.findRoot(a);
        const rootB = this.findRoot(b);

        if(rootA === rootB) return;

        if(this.catSizes[a] < this.catSizes[b]) {
            this.roots[a] = rootB;
            this.catSizes[b] += this.catSizes[a];
        }
        else {
            this.roots[b] = rootA;
            this.catSizes[a] += this.catSizes[b]
        }

        this.numCategories--;
    }

    public findRoot(x:number) {
        // find root of x
        
        // base case
        if(x === this.roots[x]) {
            return x;
        }
                
        this.roots[x] = this.findRoot(this.roots[x]);
        
        
        return this.roots[x];
    }

    public getNumCategories() {
        // return number of categories

        return this.numCategories;
    }
}