Submission #44350

#	Time	Username	Problem	Language	Result	Execution time	Memory
44350		chonka	Construction of Highway (JOI18_construction)	C++98	0 / 100	8 ms	3468 KiB

construction

#include<iostream>
#include<stdio.h>
#include<vector>
#include<map>
#include<algorithm>
using namespace std ;

#define MAXN 100007
#define LOG 21

int n ;
pair < int , int > edge[ MAXN ] ;
vector < int > v[ MAXN ] ;
int c[ MAXN ] ;
int lvl[ MAXN ] ;

int heavy[ MAXN ] ;
int subtree[ MAXN ] ;
int LCA[ MAXN ][ LOG ] ;

int root[ MAXN ] ;
int pos_in_tree[ MAXN ] ;

map < int , int > ZX ;

int MXCOLOUR ;


class Tree {
    public :
    int tr[ 4 * MAXN ] ;
    int lazy[ 4 * MAXN ] ;
    void init ( int where , int IL , int IR ) {
        tr[ where ] = lazy[ where ] = 0 ;
        if ( IL == IR ) { return ; }
        int mid = ( IL + IR ) / 2 ;
        init ( 2 * where , IL , mid ) ;
        init ( 2 * where + 1 , mid + 1 , IR ) ;
    }
    void push_lazy ( int where , int IL , int IR ) {
        if ( lazy[ where ] == 0 ) { return ; }
        tr[ where ] = lazy[ where ] ;
        if ( IL != IR ) {
            lazy[ 2 * where ] = lazy[ where ] ;
            lazy[ 2 * where + 1 ] = lazy[ where ] ;
        }
        lazy[ where ] = 0 ;
        return ;
    }
    void update ( int where , int IL , int IR , int CURL , int CURR , int nwval ) {
        push_lazy ( where , IL , IR ) ;
        if ( IR < CURL || CURR < IL ) { return ; }
        if ( CURL <= IL && IR <= CURR ) {
            lazy[ where ] = nwval ;
            push_lazy ( where , IL , IR ) ;
            return ;
        }
        int mid = ( IL + IR ) / 2 ;
        update ( 2 * where , IL , mid , CURL , CURR , nwval ) ;
        update ( 2 * where + 1 , mid + 1 , IR , CURL , CURR , nwval ) ;
    }
    int query ( int where , int IL , int IR , int pos ) {
        push_lazy ( where , IL , IR ) ;
        if ( IR < pos || pos < IL ) { return -1 ; }
        if ( IL == IR ) { return tr[ where ] ; }
        int mid = ( IL + IR ) / 2 ;
        return max ( query ( 2 * where , IL , mid , pos ) , query ( 2 * where + 1 , mid + 1 , IR , pos ) ) ;
    }
};
Tree w ;


class inv_counter {
    public :
    long long tr[ 4 * MAXN ] ;
    void init ( int where , int IL , int IR ) {
        tr[ where ] = 0 ;
        if ( IL == IR ) { return ; }
        int mid = ( IL + IR ) / 2 ;
        init ( 2 * where , IL , mid ) ;
        init ( 2 * where + 1 , mid + 1 , IR ) ;
    }
    void update ( int where , int IL , int IR , int pos , int val ) {
        if ( IR < pos || pos < IL ) { return ; }
        if ( IL == IR ) { tr[ where ] += val ; return ; }
        int mid = ( IL + IR ) / 2 ;
        if ( pos <= mid ) {
            update ( 2 * where , IL , mid , pos , val ) ;
        }
        else {
            update ( 2 * where + 1 , mid + 1 , IR , pos , val ) ;
        }
        tr[ where ] = ( tr[ 2 * where ] + tr[ 2 * where + 1 ] ) ;
    }
    long long query ( int where , int IL , int IR , int CURL , int CURR ) {
        if ( CURL > CURR ) { return 0 ; }
        if ( IR < CURL || CURR < IL ) { return 0 ; }
        if ( CURL <= IL && IR <= CURR ) {
            return tr[ where ] ;
        }
        int mid = ( IL + IR ) / 2 ;
        return ( query ( 2 * where , IL , mid , CURL , CURR ) + query ( 2 * where + 1 , mid + 1 , IR , CURL , CURR ) ) ;
    }
};

inv_counter z ;

void init ( int vertex , int prv ) {
    int i ;
    int sz = v[ vertex ].size ( ) ;
    if ( prv != -1 ) {
        for ( i = 1 ; i < LOG ; i ++ ) {
            LCA[ vertex ][ i ] = LCA[ LCA[ vertex ][ i - 1 ] ][ i - 1 ] ;
        }
    }
    subtree[ vertex ] = 1 ;
    int mx , id ;
    mx = id = -1 ;
    for ( i = 0 ; i < sz ; i ++ ) {
        if ( v[ vertex ][ i ] == prv ) { continue ; }
        lvl[ v[ vertex ][ i ] ] = lvl[ vertex ] + 1 ;
        LCA[ v[ vertex ][ i ] ][ 0 ] = vertex ;
        init ( v[ vertex ][ i ] , vertex ) ;
        subtree[ vertex ] += subtree[ v[ vertex ][ i ] ] ;
        if ( mx < subtree[ v[ vertex ][ i ] ] ) {
            mx = subtree[ v[ vertex ][ i ] ] ;
            id = v[ vertex ][ i ] ;
        }
    }
    heavy[ vertex ] = id ;
}

void compress ( ) {
    vector < int > srt ;
    srt.clear ( ) ;
    int i ;
    for ( i = 1 ; i <= n ; i ++ ) {
        srt.push_back ( c[ i ] ) ;
    }
    sort ( srt.begin ( ) , srt.end ( ) ) ;
    int id = 1 ;
    ZX[ srt[ 0 ] ] = 1 ;
    for ( i = 1 ; i < n ; i ++ ) {
        if ( srt[ i ] == srt[ i - 1 ] ) { continue ; }
        id ++ ;
        ZX[ srt[ i ] ] = id ;
    }
    MXCOLOUR = id ;
    for ( i = 1 ; i <= n ; i ++ ) {
        c[ i ] = ZX[ c[ i ] ] ;
    }
}

int get_fst ( int x ) {
    int i ;
    int sr = w.query ( 1 , 1 , n , pos_in_tree[ x ] ) ;
    for ( i = LOG - 1 ; i >= 0 ; i -- ) {
        if ( LCA[ x ][ i ] != 0 ) {
            int ret = w.query ( 1 , 1 , n , pos_in_tree[ LCA[ x ][ i ] ] ) ;
            if ( ret == sr ) { x = LCA[ x ][ i ] ; }
        }
    }
    return x ;
}

void calc ( int st ) {
    int x = st ;
    vector < pair < long long , int > > aux ;
    aux.clear ( ) ;
    while ( x != 0 ) {
        int id = get_fst ( x ) ;
        aux.push_back ( make_pair ( lvl[ x ] - lvl[ id ] + 1 , w.query ( 1 , 1 , n , pos_in_tree[ id ] ) ) ) ;
        x = LCA[ id ][ 0 ] ;
    }
    int sz = aux.size ( ) ;
    long long ans = 0 ;
    int i , j ;
    for ( i = sz - 1 ; i >= 0 ; i -- ) {
        ans += aux[ i ].first * z.query ( 1 , 1 , MXCOLOUR , aux[ i ].second + 1 , MXCOLOUR ) ;
        z.update ( 1 , 1 , MXCOLOUR , aux[ i ].second , aux[ i ].first ) ;
    }
    printf ( "%lld\n" , ans ) ;
    for ( i = 0 ; i < sz ; i ++ ) {
        z.update ( 1 , 1 , MXCOLOUR , aux[ i ].second , -aux[ i ].first ) ;
    }
}

void input ( ) {
    scanf ( "%d" , &n ) ;
    int i ;
    for ( i = 1 ; i <= n ; i ++ ) {
        scanf ( "%d" , &c[ i ] ) ;
    }
    for ( i = 1 ; i < n ; i ++ ) {
        scanf ( "%d%d" , &edge[ i ].first , &edge[ i ].second ) ;
        v[ edge[ i ].first ].push_back ( edge[ i ].second ) ;
        v[ edge[ i ].second ].push_back ( edge[ i ].first ) ;
    }
    compress ( ) ;
}

void solve ( ) {
    init ( 1 , -1 ) ;
    int i ;
    int pos = 0 ;
    for ( i = 1 ; i <= n ; i ++ ) {
        if ( i == 1 || heavy[ LCA[ i ][ 0 ] ] != i ) {
            int x = i ;
            while ( x != -1 ) {
                pos ++ ;
                pos_in_tree[ x ] = pos ;
                root[ x ] = i ;
                x = heavy[ x ] ;
            }
        }
    }
    w.init ( 1 , 1 , n ) ;
    w.update ( 1 , 1 , n , pos_in_tree[ 1 ] , pos_in_tree[ 1 ] , c[ 1 ] ) ;
    z.init ( 1 , 1 , MXCOLOUR ) ;
    for ( i = 1 ; i < n ; i ++ ) {
        calc ( edge[ i ].first ) ;
        int x = edge[ i ].second ;
        while ( x != 0 ) {
            w.update ( 1 , 1 , n , pos_in_tree[ root[ x ] ] , pos_in_tree[ x ] , c[ edge[ i ].second ] ) ;
            x = LCA[ root[ x ] ][ 0 ] ;
        }
    }
}

int main ( ) {
    input ( ) ;
    solve ( ) ;
    return 0 ;
}

Compilation message (stderr)

construction.cpp: In function 'void calc(int)':
construction.cpp:177:13: warning: unused variable 'j' [-Wunused-variable]
     int i , j ;
             ^
construction.cpp: In function 'void input()':
construction.cpp:189:11: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
     scanf ( "%d" , &n ) ;
     ~~~~~~^~~~~~~~~~~~~
construction.cpp:192:15: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
         scanf ( "%d" , &c[ i ] ) ;
         ~~~~~~^~~~~~~~~~~~~~~~~~
construction.cpp:195:15: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
         scanf ( "%d%d" , &edge[ i ].first , &edge[ i ].second ) ;
         ~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

• Subtask 1: 0 / 7
• Subtask 2: 0 / 9
• Subtask 3: 0 / 84