torus711 のアレ

主に競技プログラミングの問題について書きます

TopCoder SRM 584, Division 2, level 3 : Excavations2

概要

N 個の遺跡が地中に埋もれており、それらの種類は [ 1, 50 ] の整数で区別される。
K 個の遺跡が発掘され、発見された遺跡の種類のリストが与えられる。
(同種の遺跡が複数発掘された場合もリストに数字が現れるのは一度である)
発掘された K 個の遺跡の組み合わせの数を求めよ。

解法

答えが大きい数え上げなので多分 DP だろう、ということで動的計画法で解けないか考えます。
遺跡を一つずつ処理するようにすると、どの種類の遺跡が発掘されたかを状態にする必要があって、これは 2^50 になるので大きすぎます。
では、遺跡を種類ごとにまとめて処理するとどうなるでしょうか。
この場合、各種類ごとに一つ以上選択するか、全く選ばないようにすれば found の条件を満たすことができます。
なので、dp[ 種類 i まで考慮して ][ j 個選んだ ] := 組み合わせの数 として DP できます。
初期値は dp[0][0] = 1 です。
種類ごとにいくつあるかを予め数えて nums という配列に入れたとして、状態遷移は次のようになります。

  • i + 1 が found に含まれる場合
    • dp[ i + 1 ][ j + k ] += dp[ i ][ j ] * _{nums[ i ]}\mathrm{C}_k ( ただし 1 <= k <= nums[i] )
  • それ以外の場合
    • dp[ i + 1 ][ j ] += dp[ i ][ j ]

_n\mathrm{C}_k が出てきましたが、これも DP で求めることができます。、
dp[ i 個まで考慮して ][ j 個選んだ ] := 組み合わせの数 として、初期状態は dp[0][0] = 1 です。
状態遷移は次のようになります。

  • dp[ i + 1 ][ j + 1 ] += dp[ i ][ j ]
  • dp[ i + 1 ][ j ] += dp[ i ][ j ]

dp[ n ][ k ] が _n\mathrm{C}_k の値になります。

コード

typedef long long LL;
typedef vector<int> VI;

#define REP( i, m, n ) for ( int i = (int)( m ); i < (int)( n ); ++i )
#define ALL( c ) (c).begin(), (c).end()

typedef vector<LL> VL;
typedef vector<VL> VVL;

class Excavations2
{
public:
	long long count( vector <int> kind, vector <int> found, int K )
	{
		VI nums( 50, 0 );
		REP( i, 0, kind.size() )
		{
			nums[ kind[i] - 1 ] ++;
		}
		sort( ALL( found ) );

		VVL dp( 51, VL( 51, 0 ) );
		dp[0][0] = true;

		REP( i, 0, 50 )
		{
			REP( j, 0, 50 )
			{
				if ( binary_search( ALL( found ), i + 1 ) )
				{
					REP( k, 1, nums[i] + 1 )
					{
						dp[ i + 1 ][ j + k ] += dp[i][j] * combination( nums[i], k );
					}
				}
				else
				{
					dp[ i + 1 ][j] += dp[i][j];
				}
			}
		}

		return dp.back()[K];
	}

	LL combination( LL n, LL k )
	{
		VVL dp( n + 1, VL( k + 2, 0 ) );
		dp[0][0] = 1;

		REP( i, 0, n )
		{
			REP( j, 0, k + 1 )
			{
				dp[ i + 1 ][ j + 1 ] += dp[i][j];
				dp[ i + 1 ][j] += dp[i][j];
			}
		}
		return dp[n][k];
	}
};