*[hackerrank]Similar Pair

https://www.hackerrank.com/contests/101hack/challenges/similarpair

这题有点意思。在DFS遍历的过程中，用线段树去记录和查询区间的点数。要注意线段树的这种写法里面，index和begin，end是融为一体的，索引其实就决定了begin和end。begin，end，left，right，node，都是值域空间里的。

#include <cmath>

#include <cstdio>

#include <vector>

#include <iostream>

#include <algorithm>

using namespace std;



typedef long long llong;



vector<vector<int>> graph;

vector<int> deg; // in-degree

vector<int> tree;

int n, t;



// [left, right] is the segment to search

int query(int index, int left, int right, int begin, int end) {

    if (right < begin || left > end) return 0;

    if (left <= begin && right >= end) return tree[index];

    int mid = begin + (end - begin) / 2;

    return query(index * 2, left, right, begin, mid) + query(index * 2 + 1, left, right, mid + 1, end);

}



void update(int index, int node, int val, int begin, int end) {

    if (node < begin || node > end) {

        return;

    }

    if (begin <= node && node <= end) {

        tree[index] += val;    

    }

    if (begin == node && end == node) {

        return;

    }

    int mid = begin + (end - begin) / 2;

    update(index * 2, node, val, begin, mid);

    update(index * 2 + 1, node, val, mid + 1, end);

}



llong similarPairs(int root) {

    int l = max(1, root - t);

    int r = min(n, root + t);

    llong sim = query(1, l, r, 1, n);

    update(1, root, 1, 1, n);

    for (int i = 0; i < graph[root].size(); i++) {

        sim += similarPairs(graph[root][i]);

    }

    update(1, root, -1, 1, n);

    return sim;

}



int main() {

    cin >> n >> t;

    graph.resize(n + 1);

    deg.resize(n + 1);

    tree.resize(n * 4);

    for (int i = 1; i < n; i++) {

        int a, b;

        cin >> a >> b;

        graph[a].push_back(b);

        deg[b]++;

    }

    int root = -1;

    for (int i = 1; i < deg.size(); i++) {

        if (deg[i] == 0) {

            root = i;

            break;

        }

    }

    llong result = similarPairs(root);

    cout << result << endl;

    return 0;

}