BZOJ-1468: Tree（树-点分治）

题目：http://www.lydsy.com/JudgeOnline/problem.php?id=1468

题目大意：给定一棵边带权的树，求路径长度不大于k的路径总数。

代码（树的点分治算法，具体可看09年qzc的集训队论文《分治算法在树的路径问题中的应用》O( n log^2 n )）：

#include 

#include 

#include 

 

using namespace std ;

 

#define AddEdge( s , t , d )Add( s , t , d ),Add( t , s , d )

#define MAXN 40010

#define L( t ) left[ t ]

#define K( t ) key[ t ]

#define R( t ) right[ t ]

#define S( t ) size[ t ]

#define MAXV 100100

#define ll long long

 

struct edge{

    edge *next ;

    int t , d ;

}*head[ MAXN ];

 

void Add(int s ,int t ,int d ){

    edge *p =new( edge );

    p -> t = t , p -> d = d , p -> next = head[ s ];

    head[ s ]= p ;

}

 

int n , w ;

 

struct SBT{

    int left[ MAXV ], right[ MAXV ], key[ MAXV ], size[ MAXV ], V , roof ;

    SBT(  ){

        L(0)=R(0)=K(0)=S(0)=0;

    }

    void Init(  ){

        roof = V =0;

    }

    void update(int t ){

        S( t )=S(L( t ))+S(R( t ))+1;

    }

    void Left(int&t ){

        int k =R( t );

        R( t )=L( k );update( t );

        L( k )= t ;update( k );

        t = k ;

    }

    void Right(int&t ){

        int k =L( t );

        L( t )=R( k );update( t );

        R( k )= t ;update( k );

        t = k ;

    }

    void maintain(int&t ){

        if(S(L(L( t )))>S(R( t ))){

            Right( t );

            maintain(R( t ));maintain( t );

            return;

        }

        if(S(R(L( t )))>S(R( t ))){

            Left(L( t ));Right( t );

            maintain(L( t )),maintain(R( t ));maintain( t );

            return;

        }

        if(S(R(R( t )))>S(L( t ))){

            Left( t );

            maintain(L( t ));maintain( t );

            return;

        }

        if(S(L(R( t )))>S(L( t ))){

            Right(R( t ));Left( t );

            maintain(L( t )),maintain(R( t ));maintain( t );

            return;

        }

    }

    void Insert(int k ,int&t ){

        if(! t ){

            t =++ V ;

            K( t )= k ,L( t )=R( t )=0,S( t )=1;

            return;

        }

        Insert( k , k  next )if( vis[ p -> t ]&& p -> t != u ){

        dfs0( p -> t , v );

        size[ v ]+= size[ p -> t ];

    }

}

 

void dfs1(int v ,int u ){

    if( roof )return;

    bool flag =true;

    if( size[ rt ]- size[ v ]> size[ rt ]/2) flag =false;

    for( edge *p = head[ v ]; p ; p = p -> next )if( vis[ p -> t ]&& p -> t != u ){

        if( size[ p -> t ]> size[ rt ]/2) flag =false;

        dfs1( p -> t , v  );

    }

    if( flag ) roof = v ;

}

 

ll dep[ MAXN ];

 

void dfs2(int v ,int u ){

    for( edge *p = head[ v ]; p ; p = p -> next )if( vis[ p -> t ]&& p -> t != u ){

        dep[ p -> t ]= dep[ v ]+( ll )( p -> d );

        dfs2( p -> t , v );

    }

}

 

int a[ MAXN ], an ;

 

void dfs3(int v ,int u ){

    a[++ an ]= v ;

    for( edge *p = head[ v ]; p ; p = p -> next )if( vis[ p -> t ]&& p -> t != u ){

        dfs3( p -> t , v );

    }

}

 

void Solve(int v ){

    sbt.Init(  );

    dfs0( v ,0);

    rt = v , roof =0;

    dfs1( v ,0);

    dep[ roof ]=0;

    dfs2( roof ,0);

    sbt.Insert(0, sbt.roof );

    for( edge *p = head[ roof ]; p ; p = p -> next )if( vis[ p -> t ]){

        an =0;

        dfs3( p -> t , roof );

        for(int i =0; i ++< an ;) ans +=( ll )( sbt.Rank( w - dep[ a[ i ]], sbt.roof ));

        for(int i =0; i ++< an ;) sbt.Insert( dep[ a[ i ]], sbt.roof );

    }

    vis[ roof ]=false;

    for( edge *p = head[ roof ]; p ; p = p -> next )if( vis[ p -> t ]){

        Solve( p -> t );

    }

}

 

int main(  ){

    memset( head ,0,sizeof( head ));

    scanf("%d",&n );

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

        int s , t , d ;scanf("%d%d%d",&s ,&t ,&d );

        AddEdge( s , t , d );

    }

    scanf("%d",&w );

    memset( vis ,true,sizeof( vis ));

    Solve(1);

    printf("%lld\n", ans );

    return 0;

}