Problem Statement
You are given a tree where each node is labeled from 1 to n. How many similar pairs(S) are there in this tree?
A pair (A,B) is a similar pair if the following are true:
- node A is the ancestor of node B
- abs(A−B)≤T
Input format:
The first line of the input contains two integers, n and T. This is followed by n−1 lines, each containing two integers si and ei where node si is a parent to node ei.
Output format:
Output a single integer which denotes the number of similar pairs in the tree.
Constraints:
1≤n≤100000
0≤T≤n
1≤si, ei ≤n
Sample Input:
5 2
3 2
3 1
1 4
1 5
Sample Output:
4
Explanation:
The similar pairs are: (3, 2) (3, 1) (3, 4) (3, 5).
You can have a look at the tree image here
#include <bits/stdc++.h>
using namespace std;
int T, n;
const int N(1e5+);
int bit[N];
int sum(int x){
int s=;
while(x){
s+=bit[x];
x-=x&-x;
}
return s; //error-prone
}
void add(int x){
while(x<=n){
bit[x]++;
x+=x&-x;
}
}
int get_ans(int x){
int l=max(x-T-, );
int r=min(n, x+T);
return sum(r)-sum(l);
}
int par[N];
vector<int> g[N];
long long ans;
void dfs(int u){
int tmp=get_ans(u);
for(int i=; i<g[u].size(); i++){
int &v=g[u][i];
dfs(v);
}
ans+=get_ans(u)-tmp;
add(u);
}
int main(){
//freopen("in", "r", stdin);
cin>>n>>T;
for(int i=, u, v; i<n; i++){
cin>>u>>v;
par[v]=u;
g[u].push_back(v);
}
int root=;
while(par[root])
root=par[root];
dfs(root);
cout<<ans<<endl;
}
--------------------------------------------
我们在考虑能否用C++ STL中的 set实现这个集合
显然我们需要支持3种操作
1. 插入，set OK
2.查询集合中大于x的数有多少个
3.查询集合中小于x的数有多少个
下面的资料摘自C++ Primer (5th. edition)
P. 330
Table 9.2 Contianer Operations
Type Aliases
difference_type Signed integral type big enough to hold the distance between two iterators
set<int> s;
int f(int l, int r){
return s.upper_bound(r)-s.lower_bound(l);
}
但这是行不通的，编译时报错：
：no match for ‘operator-’ (operand types are ‘std::set<int>: :iterator {aka std::_Rb_tree_const_iterator<int>}’ and ‘std::set<int>::iterator {aka std::_Rb_tree_const_iterator<int>}’)
因为set<int>::iterator不支持-（减法）
------------------------------------------------
只能写成
set<int> s;
int f(int l, int r){
auto b=s.lower_bound(l), e=s.upper_bound(r);
int res=;
while(b!=e){ //use != rather than <
++b;
++res;
}
return res;
}
但这样写复杂度是O(n)，不能承受。
Bjarne Stroustrup TC++PL (4th. edition) P.954
The reason to use != rather than < for testing whether we have reached the end is partially because that is
the more precise statement of what we testing for and partially because only random-access iterators support <.
----------------------------------------------------------