官方题解｜二叉树叶结点的深度之和

题目解析

本题考查树的重构，根据二叉树的先序遍历和中序遍历，构建二叉树。

由于题目中给出的先序遍历和中序遍历不一定合法，所以我们首先思考，合法的先序遍历和中序遍历是什么样子的。

对于一个颗二叉树的先序遍历和中序遍历：

1. 令当前先序序列为 $[\tt{preBegin}, \tt{preEnd})$；
2. 令当前中序序列为 $[\tt{inBegin}, \tt{inEnd})$；
3. 当前先序序列的第一个元素即 $\tt{preBegin}$ 上的元素，为当前子序列对应子树的根节点；
4. 在当前中序序列中找到根节点的位置 $\tt{rootOfInOrd}$，访问当前子树根节点；
5. 将当前中序序列以 $\tt{rootOfInOrd}$ 分为左子树 $[\tt{inBegin},\ \tt{rootOfInOrd})$ 和 右子树 $[\tt{rootOfInOrd+1},\ \tt{inEnd})$ 两部分；同时可以获取左子树的大小 $\tt{loLen} = \tt{rootOfInOrd} - \tt{inBegin}$ 和右子树的大小 $\tt{hiLen} = \tt{inEnd} - \tt{(rootOfInOrd + 1)}$；
6. 根据左子树的大小和右子树的大小，将先序序列分为左子树 $[\tt{preBegin+1},\ \tt{preBegin+1+loLen})$ 和右子树 $[\tt{preBegin+1+loLen},\ \tt{preEnd})$；
7. 重复以上步骤，直至先序序列和中序序列为空。

上述过程中，若发现当前的先序序列中元素和中序序列中的元素不匹配则说明序列不合法。我们可以使用集合 std::set 来完成这个判断。

对于本题我们需要访问所有叶结点，并统计叶结点的深度信息。
那么在构建过程中，带着参数 $\tt{depth}$ 记录这一信息即可。若发现当前先序序列和中序序列的大小为 $1$，则说明序列中元素为叶结点。在递归的过程中，递归获取左子树的叶结点深度之和 $\tt{loSum}$，右子树的叶结点深度之和 $\tt{hiSum}$，$\tt{loSum + hiSum}$ 便是当前子树中所有叶结点深度之和。

AC代码

#include <bits/stdc++.h>
using namespace std;

using Iter = vector<int>::iterator;

bool checkValid(Iter preBegin, Iter preEnd, Iter inBegin, Iter inEnd) {
    set<int> st;
    while (preBegin != preEnd)
        st.insert(*preBegin++);
    while (inBegin != inEnd)
        if (!st.count(*inBegin++))
            return false;
    return true;
}

int dfs(Iter preBegin, Iter preEnd, Iter inBegin, Iter inEnd, int depth) {
    if (!checkValid(preBegin, preEnd, inBegin, inEnd)) return -1;
    if (inBegin == inEnd) return 0;
    if (distance(inBegin, inEnd) == 1) return depth;
    Iter rootOfInOrd = find(inBegin, inEnd, *preBegin);
    int loLen = distance(inBegin, rootOfInOrd);
    int hiLen = distance(rootOfInOrd+1, inEnd);
    int loSum = dfs(preBegin+1, preBegin+1+loLen, inBegin, rootOfInOrd, depth+1);
    int hiSum = dfs(preBegin+1+loLen, preEnd, rootOfInOrd+1, inEnd, depth+1);
    if (loSum == -1 or hiSum == -1) return -1;
    return loSum + hiSum;
}

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int _; cin >> _;
    while (_--) {
        int n; cin >> n;
        vector<int> a(n), b(n);
        for (auto &x : a) cin >> x;
        for (auto &x : b) cin >> x;
        cout << dfs(begin(a), end(a), begin(b), end(b), 0) << '\n';
    }
    return 0;
}

复杂度分析

根据先序遍历和中序遍历构建二叉树的过程中，朴素方法寻找根节点的总时间复杂度为 $\mathrm{O}(N^2)$；
判断先序序列中元素和中序序列中的元素是否匹配的时间复杂度为 $\mathrm{O}(N\log{N})$；
整个算法时间复杂度为 $\mathrm{O}(N^2)$。