Untitled

 avatar
Darin
plain_text
2 years ago
2.9 kB
4
Indexable
import java.util.*;
 
 
class UserSolution {
    Map<Integer, Node> map = new HashMap<>();
    Node root;
    int curGroup;
 
 
    void init(int mId, int mNum) {
        root = new Node(mId, mNum);
        map.put(mId, root);
    }
 
 
    int add(int mId, int mNum, int mParent) {
        Node parent = map.get(mParent);
        if (parent.child.size() == 2)
            return -1; // p already has 2 child nodes
        Node node = new Node(mId, mNum);
        map.put(mId, node);
        // connect n and parent p
        node.parent = parent;
        parent.child.add(node);
 
 
        Node cur = parent;
        while (cur != null) { // recursive to top to add mNum into every ascendant node
            cur.total += mNum;
            cur = cur.parent;
        }
        return parent.total;
    }
 
 
    int remove(int mId) {
        Node node = map.remove(mId);
        if (node == null)
            return -1;
        node.isRemoved = true;
        Node parent = node.parent;
        parent.child.remove(node); // remove n from its parent
 
 
        Node cur = parent;
        while (cur != null) { // recursive to top to subtract node. all People from every ascendant node
            if (cur.isRemoved == true)
                return -1;
            cur.total -= node.total;
            cur = cur.parent;
        }
        return node.total;
    }
 
 
    int tryCut(Node node, int M, int K) {
        if (node.num > K || curGroup > M) {
            curGroup = M + 1;
            return 0;
        }
        int numL, numR;
        numL = numR = 0;
 
 
        int size = node.child.size();
        if (size >= 1)
            numL = tryCut(node.child.get(0), M, K);
        if (size >= 2)
            numR = tryCut(node.child.get(1), M, K);
 
 
        if (node.num + numL > K) { // cut the left subtree
            numL = 0;
            curGroup++;
        }
        if (node.num + numR > K) { // cut the right subtree
            numR = 0;
            curGroup++;
        }
        if (node.num + numL + numR > K) { // cut the bigger subtree
            if (numL > numR)
                numL = 0;
            else
                numR = 0;
            curGroup++;
        }
        // return the remain people for current subtree
        return node.num + numL + numR;
    }
 
 
    int reorganize(int M, int K) {
        curGroup = 1;
        tryCut(root, M, K);
        return curGroup <= M ? 1 : 0;
    }
}
 
 
class Node { // represent for a department
    int id;
    int num; // people in current department
    int total; // people in the subtree
    boolean isRemoved;
    Node parent;
    List<Node> child = new ArrayList<>();
 
 
    Node(int mId, int mNum) {
        id = mId;
        num = total = mNum;
    }
}
Editor is loading...