Untitled

#define MAX_NODE 18001
#define HASH_SIZE 36001

using namespace std;

struct Node {
	int mId;
    Node* parent;
    Node* child[3];
    int num;
	int hashNext;
    bool valid;

	Node(){
		parent = nullptr;
		for(int i = 0; i < 3; i++) child[i] = nullptr;
		num = 0;
		hashNext = -1;
		valid = false;
	}

	bool addChild(Node* n){
		for(int i = 0; i < 3; i++){
			if(child[i] == nullptr){
				child[i] = n;
				return true;
			}
		}
		return false;
	}

	void removeChild(Node* n){
		for(int i = 0; i < 3; i++){
			if(child[i] == n){
				child[i] = nullptr;
			}
		}
	}
}departments[MAX_NODE];

int departmentCnt, groupCnt;
int hashTb[HASH_SIZE];

void addHash(int id)
{
	int h = departments[id].mId % HASH_SIZE;
	departments[id].hashNext = hashTb[h];
	hashTb[h] = id;
}

int getHash(int mId)
{
	int h = mId % HASH_SIZE;
	int id = hashTb[h];
	while(id != -1 && departments[id].mId != mId){ 
		id = departments[id].hashNext;
	}
	return id;
}

void init(int N, int mId[], int mNum[]) {
    departmentCnt = 0;
	for(int i = 0; i < HASH_SIZE; i++){
		hashTb[i] = -1;
	}
	for(int i = 0; i < MAX_NODE; i++){
		departments[i] = Node();
	}

    groupCnt = N;
    for (int i = 0; i < N; ++i) {
        Node* node = &departments[departmentCnt];
        node->valid = true;
        node->num = mNum[i];
		node->mId = mId[i];
		addHash(departmentCnt);

		departmentCnt++;
    }
}


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

	if(!parent->addChild(node))
		return -1;

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

	node->mId = mId;
    node->valid = true;
    node->parent = parent;
    node->num = mNum;
	addHash(departmentCnt);

    departmentCnt++;
    return parent->num;
}


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

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

	node->valid = false;
	parent->removeChild(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;
}