SRM537 Div2 Hard(925) PrinceXToastbook

PrinceXToastbook

ノード i から、prerequisite[i] に枝を張ると森になる。その森を辿って、あるトーストを p 番目に食べたときに、そのトーストと子孫のうち何枚を学習できるかを求める。

#include <vector>
#include <numeric>
using namespace std;

vector<int> P;

//  配列Sを返す
//  S[i]は、p番目のトーストをi番目に食べた場合の、
//  p番目のトーストとその子孫で学習できる枚数の期待値
//  cはルートからp番目のトーストまでの深さ
vector<double> solve( int p, int c )
{
    int n = (int)P.size();

    vector<double> S(n,1);

    for ( int i=0; i<n; i++ )
    if ( P[i]==p )
    {
        vector<double> T = solve(i,c+1);
        for ( int j=0; j<n; j++ )
        for ( int k=j+1; k<n; k++ )
            S[j] += T[k]/(n-c);
    }

    return S;
}

class PrinceXToastbook{public:
double eat( vector <int> prerequisite )
{
    P = prerequisite;
    int n = (int)P.size();

    double ans = 0;
    for ( int i=0; i<n; i++ )
    if ( prerequisite[i]==-1 )
    {
        vector<double> T = solve(i,1);
        ans += accumulate(T.begin(),T.end(),0.0)/n;
    }

    return ans;
}};