题解 | #删层子树#

牛客不保存代码，记录一下

数据结构使用queue，方便层序遍历。

class Solution {
  public:
	vector<TreeNode*> deleteLevel(TreeNode* root, vector<int>& a) {
		vector<TreeNode*> vec;
		if (root == nullptr)
			return vec;
		queue<TreeNode*> que;
        a.push_back(0);		//将root前的一层设置为要删除
		que.push(root);
		int ceng = 0;		//记录层数
		while (!que.empty()) {
			int size = que.size(); //一定要设置size，而不是动态获取（i < que.size()）,不然后面push之后，队列的大小就变了，影响ceng的大小
			ceng++;
			bool preflag = (find(a.begin(), a.end(), ceng - 1) != a.end());		//删除的是上一层
			bool nowflag = (find(a.begin(), a.end(), ceng) != a.end());			//删除的是当前层
			bool nextflag = (find(a.begin(), a.end(), ceng + 1) != a.end());	//删除的是下一层
			for (int i = 0; i < size; ++i) {
				TreeNode* node = que.front();
				que.pop();
				if (node->left) que.push(node->left);
				if (node->right) que.push(node->right);
				if (preflag) {
					if (!nowflag) {
						vec.push_back(node);
					}
				}
				if (nextflag) {
					node->left = nullptr;
					node->right = nullptr;
				}
			}
		}
		return vec;
	}
};

使用队列进行层序遍历

1. 在root节点前设置为第0层，要删除，避免根节点分情况讨论（这点很重要）。
2. 将当前节点左右子树填入遍历队列。
3. 分情况讨论

• 若当前节点上一层应该删除且这一层不用删除，将该节点加入目标数组。
• 若当前节点下一层需要删除，将当前节点的子树设置为空。