Poj Solution 2054

http://poj.org/problem?id=2054

#include <stdio.h>
#include <vector>
#include <set>
using namespace std;

struct node
{
    int id;
    int num;
    int cost;
    int sigma_c;
    node *father;
    vector <node *> child;
    bool colored;

}t[1000];

struct cmp
{
    bool operator()( node *a, node *b )const
    {
        return a->sigma_c*b->num > b->sigma_c*a->num || ( a->sigma_c*b->num == b->sigma_c*a->num && a->id < b->id );
    }
};

multiset < node*, cmp > s;

int n, r;

bool init()
{
    int i, a, b;
    node nd;

    s.clear();

    scanf( "%d%d", &n, &r );
    if( n == 0 && r == 0 ) return false;
    
    r--;

    for( i=0; i<n; i++ )
    {
        scanf( "%d", &t[i].cost );
        
        t[i].id = i;
        t[i].num = 1;
        t[i].sigma_c = t[i].cost;
        t[i].colored = false;
        t[i].father = 0;
        t[i].child.clear();

        if( i != r ) s.insert( &t[i] );
    
    }

    for( i=0; i<n-1; i++ )
    {
        scanf( "%d%d", &a, &b );
        a--, b--;

        t[a].child.push_back( &t[b] );
        t[b].father = &t[a];
    }
    
    t[r].colored = true;

    return true;
}

void doit()
{
    int k, i;
    int ans;
    node *p, *f;
    multiset < node*, cmp >::iterator iter;

    k = 1, ans = t[r].cost;

    while( !s.empty() )
    {
        iter = s.begin();
        p = *iter;
        s.erase( iter );

        if( p->father->colored )
        {
            ans += k * p->sigma_c + p->cost ;
            p->colored = true;
            k += p->num;
            continue;
        }
        
        iter = s.find( p->father );
        f = *iter;
        s.erase( iter );

        f->cost += p->cost + p->sigma_c * f->num ; 
        f->num += p->num;
        f->sigma_c += p->sigma_c;
        
        for( i=0; i<p->child.size(); i++ )
            p->child[i]->father = f;
        
        f->child.insert( f->child.end(), p->child.begin(), p->child.end() );

        p->child.clear();

        s.insert( f );
    }

    printf( "%dn", ans );
}

int main()
{
    while( init() )
        doit();
    return 0;
}