public class Solution {
	
	class TreeNode {
		int val;
		TreeNode left;
		TreeNode right;
		
	}
	
	// 用于保存节点记录
	class Record{
		TreeNode node;
		int nodeCount;
		public Record(TreeNode node, int nodeCount) {
			this.node = node;
			this.nodeCount = nodeCount;
		}
	}

	List<Record> records;
	
	int nodes;

	public TreeNode[] delete(TreeNode root) {
		if (root == null) {
			return null;
		}
		records = new ArrayList<>();
		nodes = 0;
		dfs(root, 0);
		records.add(new Record(root, nodes));
		
		Collections.sort(records, new Comparator<Record>() {
			@Override
			public int compare(Record o1, Record o2) {
				int dif = o1.nodeCount - o2.nodeCount;
				if (dif != 0) {
					return dif;	
				}
				return o1.node.val - o2.node.val;
			}
		});
		TreeNode[] res = new TreeNode[records.size()];
		int i = 0;
		for (Record record : records) {
			res[i++] = record.node;
		}
		return res;
	}

	private int dfs(TreeNode root) {
		TreeNode left = root.left;
		TreeNode right = root.right;
		int leftDep = left != null ? dfs(left) : 0;
		int leftNodeCount = nodes;
		nodes = 0;
		int rightDep = right != null ? dfs(right) : 0;
		int rightNodeCount = nodes;
		
		int dif;
		while ((Math.abs((dif = leftDep - rightDep))) > 1) {
			// 删除一棵子树时, 更新该子树的高度, 并清除节点计数
			if (dif > 1) {
				records.add(new Record(left, leftNodeCount));
				leftNodeCount = 0;
				leftDep = 0;
				root.left = null;
			} else {
				records.add(new Record(right, rightNodeCount));
				rightNodeCount = 0;
				rightDep = 0;
				root.right = null;
			}
		}
		// 每次递归, 利用 nodes 保存当前节点计数
		nodes = leftNodeCount + rightNodeCount + 1;
		return Math.max(leftDep, rightDep) + 1;
	}
}