求二叉树的带权路径长度（C语言/C++）

目录

一、题目

样例

数据范围

二、思路

三、解题方法

（一）先序遍历

实现代码

（二）层次遍历

实现代码

---

一、题目

3766. 二叉树的带权路径长度 - AcWing题库

（2014年408数据结构考题）

二叉树的带权路径长度(WPL)是二叉树中所有叶结点的带权路径长度之和，也就是每个叶结点的深度与权值之积的总和。

给定一棵二叉树 T，请你计算并输出它的 WPL。

注意，根节点的深度为 0。

样例

数据范围

二叉树结点数量不超过 1000。
每个结点的权值均为不超过 100 的非负整数。

二、思路

某叶子结点的WPL = 叶子结点权重weight * 根结点到该叶子结点路径长度depth

整棵树的WPL即为所有叶结点的带权路径长度之和。

根据公式关系，我们需要知道2件事：1.该结点是否是叶子结点；2.结点的路径长度，即该结点所处的深度。

所以，我们可以选择用先序遍历或层次遍历来实现。

三、解题方法

（一）先序遍历

1. 用一个全局变量记录wpl并初始化，把每个结点的深度作为递归函数的一个参数传递。

2. 若该结点是叶子结点，那么wpl加上该结点的深度与权值之积。

3. 若不是叶子结点，若左子树非空，对左子树调用递归算法；若右子树非空，对右子树调用递归算法。深度参数均为本结点的深度参数加1。

4. 最后返回计算出的wpl即可。

实现代码

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    int wpl = 0; //全局变量
    
    void pathOne(TreeNode* root, int deep) { //递归函数
        if(root->left == NULL && root->right == NULL) //叶子结点
            wpl += root->val*deep;
        if(root->left) { //非叶子结点，左子树非空
            pathOne(root->left, deep+1);
        }
        if(root->right) { //非叶子结点，右子树非空
            pathOne(root->right, deep+1);
        }
    }
    
    int pathSum(TreeNode* root) {
        pathOne(root, 0);
        return wpl;
    }
};

或者使用static变量记录wpl也可以。static是一个静态变量，只有在首次调用函数时声明wpl并赋值为0，以后的递归调用并不会使得wpl为0。与上面类似，实现代码如下：

class Solution {
public:
    int WPL(TreeNode* root, int depth) {
        static int wpl = 0;
        if(root->left == NULL && root->right == NULL) 
            return wpl += root->val*depth;
        if(root->left != NULL) 
            WPL(root->left, depth+1);
        if(root->right != NULL)
            WPL(root->right, depth+1);
        return wpl;
    }
    int pathSum(TreeNode* root) {
        return WPL(root, 0);
    }
};

也可以选择分别计算左子树的WPL和右子树的WPL，然后相加就是总的WPL。

class Solution {
public:
    int WPL(TreeNode* root, int depth) {
        int lwpl = 0;
        int rwpl = 0; //用于存储左子树和右子树产生的WPL
        if(root->left == NULL && root->right == NULL) //若为叶子结点，计算当前叶子结点的WPL
            return root->val*depth;
        if(root->left != NULL) //若左子树不空，对左子树递归遍历
            lwpl = WPL(root->left, depth+1);
        if(root->right != NULL) //若右子树不空，对左子树递归遍历
            rwpl = WPL(root->right, depth+1);
        return lwpl + rwpl;
    }
    int pathSum(TreeNode* root) {
        return WPL(root, 0);
    }
};

上面这一段代码，C/C++语言基础好的也可以使用更简便的以下形式：

class Solution {
public:
    int WPL(TreeNode* root, int depth) {
        if(root->left == NULL && root->right == NULL) //若为叶子结点，计算当前叶子结点的WPL
            return root->val*depth;
        return (root->left != NULL ? WPL(root->left, depth+1) : 0) + (root->right != NULL ? WPL(root->right, depth+1) : 0);
    }
    int pathSum(TreeNode* root) {
        return WPL(root, 0);
    }
};

在写代码过程种容易忘记三元式(x?y:z)两端的括号，不加括号，答案就会是错误的，切记。

（二）层次遍历

1. 使用队列进行层次遍历，并记录当前层数。

2. 当遍历到叶子结点时，累计WPL。

3. 当遍历到非叶子结点时把该结点的子树加入队列。

4. 当某结点为该层的最后一个结点，层数自增1。

5. 队列空时，返回WPL。

实现代码

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    int pathSum(TreeNode* root) {
        queue<TreeNode*> q; //声明队列
        TreeNode* lastNode = root; //lastNode用来记录当前层的最后一个结点，初始化为根结点
        TreeNode* newlastNode = NULL; //newlastNode用来记录下一层的最后一个结点，初始化为空
        int wpl = 0; //初始化wpl
        int deep = 0; //初始化深度
        q.push(root); //根结点入队
        while(!q.empty()) { //层次遍历，若队列不空则循环
            TreeNode* node = q.front(); //取出队列的头一个元素
            if(node->left == NULL && node->right == NULL) //若为叶子结点，统计wpl
                wpl += node->val * deep;
            if(node->left) { //若非叶子结点，把左结点入队
                q.push(node->left);
                newlastNode = node->left; //并设下一层的最后一个结点为该结点的左结点
            }
            if(node->right) { //若非叶子结点，把右结点入队
                q.push(node->right);
                newlastNode = node->right; //并设下一层的最后一个结点为该结点的右结点
            }
            
            if(node == lastNode) { //若该结点为本层最后一个结点，更新lastNode
                deep++; //层数+1
                lastNode = newlastNode;
            }
            q.pop(); //删除头一个元素
        }
        return wpl; //返回wpl
    }
};

C语言也可以使用数组构造队列，但需要提前设置最大容量，并且加入队列判满操作，这里队列的判满条件为front==(rear+1)%Size。