我们要知道一棵树,有几个节点。假设我们位于当前节点上,那么它需要问左右子树什么信息才能知道以当前节点为根节点的数共有几个节点呢?
它必须知道左子树,和右子树有多少节点。
因此,我们定义一个info:
class Info{
public:
int count; // 每个节点都需要知道分别向它的左右子树询问它的节点个数
explicit Info(int cnt) : count(cnt) {
}
};
接下来就是递归的套路了:
- base: 当当前节点为null时,返回Info(0);也就是没有节点
- 如果当前节点不为null,那么收集它的左右节点的Info
- 对于当前节点,根据上一步,得到它自己的Info
Info process(TreeNode* root){
if(root == nullptr){
return Info(0);
}
Info leftInfo = process(root->left);
Info rightInfo = process(root->right);
return Info(leftInfo.count + rightInfo.count + 1);
}
整体代码如下
class Solution {
struct Info{
public:
int count;
explicit Info(int cnt) : count(cnt) {
}
};
Info * process(TreeNode* root){
if(root == nullptr){
return new Info(0);
}
auto left = process(root->left);
auto right = process(root->right);
return new Info(left->count + left->count + 1);
}
public:
int countNodes(TreeNode* root) {
if(root == nullptr){
return 0;
}
return process(root)->count;
}
};