PAT 1115 Counting Nodes in a BST

典型的深度优先搜索算法题目.

在深度优先搜索过程中,记录下是否树的最大深度和每个深度上节点的个数就可以了.

#include <iostream>
#include <map>
int MAX_DEPTH=0;
using namespace std;
struct TreeNode{
    int val;
    TreeNode* left;
    TreeNode* right;
};
TreeNode* insert(TreeNode* root,int val){
    if(root==nullptr){
        root=new TreeNode();
        root->val=val;
    }else if(root->val>=val){
        root->right=insert(root->right,val);
    } else {
        root->left = insert(root->left, val);
    }
    return root;
}
map<int,int> depth_nums;
void dfs(TreeNode* node,int depth){
    if(node==nullptr){
        return;
    }
    if(depth>MAX_DEPTH){
        MAX_DEPTH=depth;
    }
    depth_nums[depth]++;
    dfs(node->left,depth+1);
    dfs(node->right,depth+1);
}
int main() {
    TreeNode* root= nullptr;
    int N;
    cin>>N;
    for (int i = 0; i < N; ++i) {
        int val;
        cin>>val;
        root=insert(root,val);
    }
    dfs(root,0);
    int num1,num2;
    num1=depth_nums[MAX_DEPTH];
    num2=depth_nums[MAX_DEPTH-1];
    cout<<num1<<" + "<<num2<<" = "<<num1+num2<<endl;
    return 0;
}