SRM476 Div1 Medium(550) FriendTour

FriendTour

Manglisiの全員を辿るのではなく、Manaoの友人をのみを辿る。friends[0]の長さが36文字以内なのでMasaoの友人は高々15人。探索の結果をメモしておけば全探索が間に合う。

d本のエッジを持つノードでそれぞれのエッジに進んだときに完遂する確率が分かっているとする。Masaoがi番目に小さな確率のエッジeを選ぶのは、表示された友人リストのなかでeが最も大きい場合。その確率は、i-1Ck-1/dCk

#include <string>
#include <vector>
#include <sstream>
#include <algorithm>
#include <map>
#include <utility>

using namespace std;

class FriendTour
{
    int n;
    vector<vector<bool> > edge;
    vector<int> degree;
    int K;

    vector<vector<int> > comb;
    map<pair<int,long long>,double> memo;

    double search( int d, int cur, long long visit );
public:
    double tourProbability( vector <string> friends, int K );
};

double FriendTour::tourProbability( vector <string> friends, int K )
{
    //  グラフを作成
    n = (int)friends.size();
    edge = vector<vector<bool> >( n, vector<bool>( n ) );
    degree = vector<int>( n );
    for ( int i=0; i<n; i++ )
    {
        stringstream ss( friends[i] );
        int f;
        while ( ss >> f )
            edge[i][f-1] = true,
            degree[i]++;
    }
    this->K = K;

    //  コンビネーションを計算
    comb = vector<vector<int> >( n+1 );
    for ( int i=0; i<=n; i++ )
    {
        for ( int j=0; j<=i; j++ )
            if ( j == 0  ||  j == i )
                comb[i].push_back( 1 );
            else
                comb[i].push_back( comb[i-1][j-1] + comb[i-1][j] );
    }

    memo.clear();

    return search( 0, 0, 0 );
}

double FriendTour::search( int d, int cur, long long visit )
{
    //  すべての友人を辿ったら終了
    if ( d == degree[0] )
        return 1;

    if ( memo.count(make_pair(cur,visit)) == 1 )
        return memo[make_pair(cur,visit)];

    //  各エッジに進んだときに完遂する確率を求める
    vector<double> prob;

    for ( int i=0; i<n; i++ )
    if ( edge[cur][i] )
        if ( edge[0][i]  &&  ( visit & 1ll<<i ) == 0 )
            prob.push_back( search(d+1,i,visit|1ll<<cur) );
        else
            prob.push_back( 0 );

    //  Manaoは表示された友達リストの中で確率が最大のものを選択する
    sort( prob.begin(), prob.end() );

    double p = 0;
    int k = min( K, degree[cur] );

    for ( int i=0; i<degree[cur]; i++ )
        if ( i >= k-1 )
            p += prob[i] * comb[i][k-1]/comb[degree[cur]][k];
    
    return memo[make_pair(cur,visit)] = p;
}