SRM474 Div1 Medium(500) TreesCount

TreesCount

ダイクストラ法で頂点vの距離を確定するときに、確定済みの頂点からvへ向かう辺のうち1本を残し他を取り除くことを考える。uとvを結ぶ辺の長さをeとして、dist[u]+e=dist[v]となる辺から残す1本を選ぶ。

#include <string>
#include <vector>

using namespace std;

class TreesCount
{
public:
    int count( vector <string> graph );
};

int TreesCount::count( vector <string> graph )
{
    int n = (int)graph.size();

    vector<int> dist( n, n*10 );
    vector<bool> fix( n, false );

    dist[0] = 0;

    long long ans = 1;

    for ( int c=0; c<n; c++ )
    {
        //  未確定かつ最短距離の頂点を探す
        int v = -1;
        for ( int i=0; i<n; i++ )
        if ( ! fix[i] )
            if ( v == -1  ||  dist[i] < dist[v] )
                v = i;

        if ( c > 0 )
        {
            //  確定済みの頂点から最短距離で
            //  この頂点に到達できる辺の本数を乗じる
            int m = 0;
            for ( int i=0; i<n; i++ )
            if ( fix[i]  &&  graph[i][v] > '0' )
                if ( dist[i]+graph[i][v]-'0' == dist[v] )
                    m++;

            ans = ans * m % 1000000007;
        }

        //  この頂点を確定
        fix[v] = true;
        for ( int i=0; i<n; i++ )
            if ( graph[v][i] > '0' )
                dist[i] = min( dist[i], dist[v]+graph[v][i]-'0' );
    }

    return (int)ans;
}