• + 0 comments

    js

    function addChild(tree, c, node, child) {
        tree[node] = tree[node] ?? { children: new Set(), val: c[node - 1] };
        tree[node].children.add(child);
    }
    
    function cleanUpTree(tree, node = 1, parent = -1, vMap = {}) {
        let { children, val } = tree[node];
        children.delete(parent);
        for (const child of children) {
            const [v] = cleanUpTree(tree, child, node, vMap);
            val += v;
        }
        vMap[val] = (vMap[val] ?? 0) + 1;
        tree[node].val = val;
        return [val, vMap];
    }
    
    function getMin(a, b) {
        return a === -1 || a > b ? b : a;
    }
    
    function* dfsIter(tree, vMap, node = 1, parentVMap = {}) {
        const { val, children } = tree[node];
        vMap[val]--;
        parentVMap[val] = (parentVMap[val] ?? 0) + 1;
        for (const child of children)
            yield* dfsIter(tree, vMap, child, parentVMap);
        parentVMap[val]--;
        if (node === 1) return;
        yield [val, parentVMap];
    }
    
    function balancedForest(c, edges) {
        let ret = -1;
        if (!c?.length || !edges?.length) return ret;
    
        const tree = {};
        for (const [begin, end] of edges) {
            addChild(tree, c, begin, end);
            addChild(tree, c, end, begin);
        }
        const [total, vMap] = cleanUpTree(tree);
    
        for (const [v1, parentVMap] of dfsIter(tree, vMap)) {
            const rest = total - v1;
    
            if (rest === v1) {
                ret = getMin(ret, v1);
                continue;
            }
    
            for (const v2 of [rest / 2, v1, rest - v1]) {
                if (!vMap[v2] && !parentVMap[v1 + v2]) continue;
                const v3 = rest - v2;
    
                const [a, b, c] = [v1, v2, v3].sort((a, b) => b - a);
                if (a !== b) continue;
    
                ret = getMin(ret, a - c);
            }
        }
        return ret;
    }