*[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; }