class Solution {
public:
    vector <TreeNode*> leafnodes;
    int count = 0;
    unordered_map <TreeNode*, unordered_map <int, TreeNode*>> parent;
    unordered_map <TreeNode*, int> dist;
    void dfs(TreeNode* source, TreeNode* p)
    {
        source->val = count;
        count++;
        parent[source][0] = p;
        int i=1;
        while(parent[source][i-1]!=NULL)
        {
            parent[source][i] = parent[parent[source][i-1]][i-1];
            i++;
        }
        if(source->left!=NULL)
        {
            dist[source->left] = dist[source] + 1;
            dfs(source->left, source);
        }
        if(source->right!=NULL)
        {
            dist[source->right] = dist[source] + 1;
            dfs(source->right, source);
        }
        if(source->left==NULL and source->right==NULL)
        {
            leafnodes.push_back(source);
        }
    }
    int LCA(TreeNode* v, TreeNode* u)
    {
        TreeNode* ov = v, *ou = u;
        if(dist[v]<dist[u])
            swap(v,u);
        int n = log(count);
        for(int i=n-1;i>=0;i--)
        {
            if(parent[v][i]!=NULL and dist[parent[v][i]]>=dist[u])
                v = parent[v][i];
        }
        if(v->val==u->val)
        {
            return dist[ov] + dist[ou] - 2*dist[v];
        }
        for(int i=n-1;i>=0;i--)
        {
            if((parent[v][i]!=NULL and parent[u][i]!=NULL) and (parent[v][i]->val != parent[u][i]->val))
            {
                v = parent[v][i], u=parent[u][i];
            }
        }
        return dist[ov] + dist[ou] - 2*dist[parent[v][0]];
    }
    int countPairs(TreeNode* root, int distance) {
        parent[root][0] = NULL;
        dist[root] = 0;
        dfs(root, NULL);
        int n = leafnodes.size();
        int count = 0;
        for(int i=0;i<n;i++)
        {
            for(int j=i+1;j<n;j++)
            {
                if(LCA(leafnodes[i], leafnodes[j])<=distance)
                    count++;
            }
        }
        return count;
    }
};