hackerrank Similar Pair

传送门

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

-------------------------------------------------------------

Solution:

题目给出一棵有根树，要求统计其中满足以下条件的点对(A, B)的数目

（1） A是B的祖先（但A不能等于B，即A应是B的合法（proper）祖先）

（2）abs(A-B)<=T（即两点编号之差的绝对值应不大于T）

--------------------------------------------------------

解法不难想，类似于树形DP，DFS遍历这棵有根树，以后序（或者说后序遍历）将各个顶点的编号逐一放到集合中。在进入以顶点u为根的子树时，先查询当前集合中有多少节点v满足 abs(v-u)<=T，遍历完这棵子树后在查询一次，由于新添进集合的都是u的后代，所以两次结果相减便是符合条件的点对(u,x)的数目。对每个节点都做类似的查询结果加起来就是答案。

--------------------------------------------------------------

至此，问题归结为如何维护这个集合，显然树状数组(BIT)是最合适的

#include <bits/stdc++.h>
using namespace std;
int T, n;
const int N(1e5+5);
int bit[N];
int sum(int x){
    int s=0;
    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-1, 0);
    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=0; 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=1, u, v; i<n; i++){
        cin>>u>>v;
        par[v]=u;
        g[u].push_back(v);
    }
    int root=1;
    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和map都支持upper_bound()，lower_bound(),可尝试用这两个函数来实现上述后两个查询。

比如我们想查询集合s中在[L, R]范围内的数有多少个，试着写成

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=0;
    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 <.

----------------------------------------------------------