• + 0 comments

    Everyone can try to solve by this js code

    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;
    

    }