UVa 1205 Color a Tree

 1  #include <stdio.h>

 2  #include <stdlib.h>

 3  #include <algorithm>

 4  #include <iostream>

 5  #include <vector>

 6  #include <string.h>

 7  using namespace std;

 8  #define MAXN 1005

 9  int n,r;

10  int cnt[MAXN],fa[MAXN],next1[MAXN],c[MAXN];

11  double now[MAXN];

12  bool vis[MAXN];

13  vector<int> edges[MAXN];

14  void ini()

15  {

16      for(int i=0; i<=n; i++)

17      {

18          cnt[i]=1;

19      }

20      memset(fa,-1,sizeof(fa));

21      memset(next1,0,sizeof(next1));

22      memset(vis,false,sizeof(vis));

23  }

24  void dfs(int x)

25  {

26      vector<int> xchild=edges[x];

27      int s=xchild.size();

28      for(int i=0; i<s; i++)

29      {

30          int j=xchild[i];

31          if(fa[j]==-1)

32          {

33              fa[j]=x;

34              dfs(j);

35          }

36      }

37  }

38  void path(int x,int y)

39  {

40      while(next1[x]) x=next1[x];

41      next1[x]=y;

42  }

43  void work()

44  {

45      double maxn=0;

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

47      {

48          int k;

49          maxn=0;

50          for(int j=1; j<=n; j++)

51          {

52              if(maxn<now[j] && vis[j]==false && j!=r)

53              {

54                  maxn=now[j];

55                  k=j;

56              }

57          }

58          int f=fa[k];

59          path(f,k);

60          while(vis[f]) f=fa[f];

61          vis[k]=true;

62          now[f]=(now[f]*cnt[f]+now[k]*cnt[k])/(cnt[f]+cnt[k]);

63          cnt[f]+=cnt[k];

64      }

65      int start=r,ans=0;

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

67      {

68          ans+=c[start]*i;

69          start=next1[start];

70      }

71      printf("%d\n",ans);

72  }

73  int main()

74  {

75      while(cin>>n>>r && r+n)

76      {

77          memset(c,0,sizeof(c));

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

79          {

80              cin>>c[i];

81              now[i]=(double)c[i];

82          }

83          for(int i=0;i<MAXN;i++) edges[i].clear();

84          for(int i=0; i<n-1; i++)

85          {

86              int a,b;

87              cin>>a>>b;

88              edges[a].push_back(b);

89          }

90          ini();

91          dfs(r);

92          work();

93      }

94      return 0;

95  }