Untitled

#include <unordered_map>
//#define MAX_NODE 18001;

using namespace std;

const int MAX_NODE = 18001;

struct Node {
    Node* parent;
    Node* child[3];
    int num;
    bool valid;
};
int departmentCnt, groupCnt;

Node departments[MAX_NODE];
unordered_map<int, int> hashMap;

void delNode(Node* node) {
    if (node == nullptr) return;

    node->valid = false;
    for (int i = 0; i < 3; ++i)
        delNode(node->child[i]);
}


void init(int N, int mId[], int mNum[]) {
    departmentCnt = 0;
    hashMap = unordered_map<int, int>();


    groupCnt = N;
    for (int i = 0; i < N; ++i) {
        hashMap[mId[i]] = departmentCnt;
        Node* node = &departments[departmentCnt++];
        node->valid = true;
        node->parent = nullptr;
        node->child[0] = node->child[1] = node->child[2] = nullptr;
        node->num = mNum[i];
    }
}


int add(int mId, int mNum, int mParent) {
    Node* node = &departments[departmentCnt];
    Node* parent = &departments[hashMap[mParent]];


    if (parent->child[0] == nullptr) {
        parent->child[0] = node;
    }
    else if (parent->child[1] == nullptr) {
        parent->child[1] = node;
    }
    else if (parent->child[2] == nullptr) {
        parent->child[2] = node;
    }
    else {
        return -1;
    }


    Node* curr = parent;
    while (curr) {
        curr->num += mNum;
        curr = curr->parent;
    }

    hashMap[mId] = departmentCnt;
    node->valid = true;
    node->parent = parent;
    node->child[0] = node->child[1] = node->child[2] = nullptr;
    node->num = mNum;

    departmentCnt++;
    return parent->num;
}


int remove(int mId) {
    Node* node = &departments[hashMap[mId]];
    if (!node || !node->valid) {
        return -1;
    }

    Node* parent = node->parent;
    Node* curr = parent;
   /* while (curr) {
        if (!curr->valid)
            return -1;
        curr = curr->parent;
    }*/

    if (parent->child[0] == node) {
        parent->child[0] = nullptr;
    }
    else if (parent->child[1] == node) {
        parent->child[1] = nullptr;
    }
    else if (parent->child[2] == node) {
        parent->child[2] = nullptr;
    }

    curr = parent;
    while (curr) {
        curr->num -= node->num;
        curr = curr->parent;
    }

    delNode(node);
    return node->num;
}


int distribute(int K) {
    int low = 1, high = 0;
    for (int i = 0; i < groupCnt; ++i) {
        if (departments[i].num > high)
            high = departments[i].num;
    }

    while (low <= high) {
        int mid = (low + high) / 2;
        int sum = 0;
        for (int i = 0; i < groupCnt; ++i) {
            if (departments[i].num <= mid)
                sum += departments[i].num;
            else
                sum += mid;
        }
        if (sum <= K) {
            low = mid + 1;
        }
        else {
            high = mid - 1;
        }
    }
    return high;
}