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

说明

连续在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 children;

    Node() {}

    Node(int _val, vector _children) {
        val = _val;
        children = _children;
    }
};
*/
class Solution {
public:
    vector> levelOrder(Node* root) {
        vector> res;
        if(root == NULL)
            return res;
        int count = 0;
        int newCount = 0;
        int c = 0;
        vector vec;    // all node`val at same level
        
        queue 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

你可能感兴趣的:(树的层序遍历并统计每一层的值)