UVA 1329 - Corporative Network ( Union Find )

UVa Online Judge
Make the disjoint set remember the distance from each node to its direct father. Whenever path compression is in progress, the cost should be updated as well.

struct UnionFind{
    vector< int > fa, cost;
    UnionFind( int sz ){
        fa.resize( sz );
        cost.resize( sz );
        for( int i = 0; i < sz; ++i )
            fa[ i ] = i,
            cost[ i ] = 0;
    }
    int find( int x ){
        if( fa[ x ] == x ) return x;
        find( fa[ x ] );
        cost[ x ] += cost[ fa[ x ] ];
        return fa[ x ] = find( fa[ x ] );
    }
    bool unite( int x, int y ){
        int a = find( x );
        int b = find( y );
        if( a == b ) return false;
        fa[ a ] = b; return true;
    }
};

void solve(){
    int N; cin >> N;
    UnionFind uf = UnionFind( N );
    string op;
    while( cin >> op and op[ 0 ] != 'O' ){
        if( op[ 0 ] == 'E' ){
            int x; cin >> x;
            uf.find( x - 1 );
            cout << uf.cost[ x - 1 ] << "\n";
        } else{
            int x, y; cin >> x >> y;
            uf.fa[ x - 1 ] = y - 1;
            uf.cost[ x - 1 ] = abs( x - y ) % 1000;
        }
    }
}