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;
}