树的层序遍历并统计每一层的值

说明

连续在LeetCode中看到好几个类似的题目,核心思想就是树的广度优先搜索(BFS),并在搜索的过程中记录每一层的值。
以LeetCode中的N-ary Tree Level Order Traversal为例:

Description

Given an n-ary tree, return the level order traversal of its nodes’ values. (ie, from left to right, level by level).

For example, given a 3-ary tree:
example-img
We should return its level order traversal:

[
    [1],
    [3,2,4],
    [5,6]
]

解析

利用队列使用BFS。设res为要返回的2维数组, vec为当前访问层的元素的1维数组。

  1. count为当前层的节点数量,初始化为0,newCount为新加入的一层的节点数量,初始化为0,c为当前访问层的第几个元素,初始化为0.
  2. 根节点入队,++count, 根节点元素加入到vec中
  3. 如果c等于count,将1维数组vec加入到res中,并使 count = newCount, c = 0, newCount = 0, 再清空vec
  4. 队列中的第一个元素出队,并赋值给p,对于p中的每一个子节点,如果不为空则进队并使newCount = newCount+1
  5. 节点p的元素加入到vec中, 并让 c = c + 1
  6. 重复步骤3-5直到队列为空
  7. 如果count大于0,则将vec中加入到res中

时间复杂度: O(n) n为节点的数量

C++实现

/* // Definition for a Node. class Node { public: int val = NULL; vector<Node*> children; Node() {} Node(int _val, vector<Node*> _children) { val = _val; children = _children; } }; */ class Solution { public: vector<vector<int>> levelOrder(Node* root) { vector<vector<int>> res; if(root == NULL) return res; int count = 0; int newCount = 0; int c = 0; vector<int> vec; // all node`val at same level queue<Node*> q; q.push(root); count = 1; while(!q.empty()){ if(c == count){ res.push_back(vec); vec.clear(); count = newCount; c = 0; newCount = 0; } auto p = q.front(); q.pop(); vec.push_back(p->val); for(int i = 0; i < p->children.size(); ++i){ if(p->children[i] != NULL){ q.push(p->children[i]); ++newCount; } } ++c; } if(count){ // last leavel values res.push_back(vec); } return res; } }; // cin优化 static const auto __ = [](){ ios::sync_with_stdio(false); cin.tie(nullptr); return nullptr; }();

这个题中不得不提一下c++中cin的优化问题,优化代码如上,主要就是关闭stdio同步和解除与cout的绑定,这份代码在LeetCode提交没优化cin之前大概为比20%的提交快,而开了cin优化便几乎比100%的提交快了…
关于cin读取数据的问题见这篇文章


本作品采用知识共享署名-非商业性使用-禁止演绎 4.0 国际许可协议进行许可。转载请注明: 作者staneyffer,首发于我的博客,原文链接: https://chengfy.com/post/8


载入评论中....