Untitled

package PhanBoPhieu;

import java.util.HashMap;
import java.util.TreeSet;

class UserSolution {
	class Node {
		int Id;
		int num_employee;
		int count;
		int num_child;
		boolean check;
		Group group;
		Node parent;
		public Node(int id, int num_employee, Group group) {
			super();
			Id = id;
			this.num_employee = num_employee;
			this.group = group;
			this.parent = null;
			this.check=true;
		}
		
	}
	class Group implements Comparable<Group>{
		int id;
		int num_employee;
		@Override
		public int compareTo(Group o) {
			if(num_employee==o.num_employee){
				return id-o.id;
			}
			return o.num_employee-num_employee;
		}
		
		
	}
	TreeSet<Group> tree;
	int total;
	HashMap<Integer, Node> nodeMap;
	public void init(int N, int mId[], int mNum[]) {
		total=0;
		tree=new TreeSet<UserSolution.Group>();
		nodeMap= new HashMap<Integer, UserSolution.Node>();
		for(int i=1;i<=N;i++){
			Group group=new Group();
			group.id=i;
			group.num_employee+=mNum[i-1];
			Node node=new Node(mId[i-1], mNum[i-1], group);
			tree.add(group);
			nodeMap.put(node.Id, node);
			total+=mNum[i-1];
		}
		return;
	}

	public int add(int mId, int mNum, int mParent) {
		Node node=nodeMap.get(mParent);
		if(node==null || node.count==3 || node.check==false){
			return -1;
		}
		Node node1=new Node(mId, mNum, node.group);
		node1.parent=node;
		node.count++;
		total+=mNum;
		nodeMap.put(node1.Id, node1);
		while(node!=null){
			node.num_child+=mNum;
			node=node.parent;
			if(node!=null && node.check==false){
				return -1;
			}
		}
		tree.remove(node1.group);
		node1.group.num_employee+=mNum;
		tree.add(node1.group);
		
		return node1.parent.num_employee+node1.parent.num_child;
	}

	public int remove(int mId) {
		Node node=nodeMap.get(mId);
		if(node.check==false){
			return -1;
		}
		node.check=false;
		tree.remove(node.group);
		node.group.num_employee-=node.num_employee+node.num_child;
		tree.add(node.group);
		Node node1=node.parent;
		node1.count--;
		while(node1!=null){
			if(node1.check==false){
				return -1;
			}
			node1.num_child-=node.num_employee+node.num_child;
			node1=node1.parent;
			
		}
		total-=node.num_employee+node.num_child;
		return node.num_employee+node.num_child;
	}

	public int distribute(int K) {
		int ans=0;
		if(K>=total){
			return tree.first().num_employee;
		}
		int R=tree.first().num_employee;
		int L=tree.last().num_employee;
		while(L<=R){
			int cnt=0;
			int mid=(R+L)/2;
			for(Group group:tree){
				if(group.num_employee<mid){
					cnt+=group.num_employee;
				}
				else{
					cnt+=mid;
				}
			}
			
			if(cnt>K){
				R=mid-1;
			}
			else{
				L=mid+1;
				if(ans<mid){
					ans=mid;
				}
			}
			
		}
		return ans;
	}
}