18.两节点的路径长度

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
    int n,m,h[1111],ne[2222],p[2222],d[1111][1111];
    bool b[1111];//标记是否叶子结点
    void add(int x,int y)
    {
        p[++m]=y;
        ne[m]=h[x];
        h[x]=m;
        p[++m]=x;
        ne[m]=h[y];
        h[y]=m;
    }
    int dfs(TreeNode* root)
    {
        int r=n++,c;
        b[r]=1;
        if(root->left!=nullptr)
        {
            c=dfs(root->left);
            add(r,c);
            b[r]=0;
        }
        if(root->right!=nullptr)
        {
            c=dfs(root->right);
            add(r,c);
            b[r]=0;
        }
        return r;
    }
    void work(int x,int r,int f)
    {
        for(int i=h[x];i;i=ne[i])if(p[i]!=f)
        {
            d[r][p[i]]=d[r][x]+1;
            work(p[i],r,x);
        }
    }
public:
    int countPairs(TreeNode* root, int distance) {
        n=m=0;
        memset(h,0,sizeof(h));
        memset(d,0,sizeof(d));
        memset(b,0,sizeof(b));
        dfs(root);
        int i,j,ans;
        for(i=0;i<n;i++)work(i,i,i);
        for(i=0;i<n;i++)if(b[i])for(j=0;j<i;j++)if(b[j]&&d[i][j]<=distance)ans++;
        return ans;
    }
};