TopCoder, SRM 619, Division 2, Level 1 : GoodCompanyDivTwo

問題文

http://community.topcoder.com/stat?c=problem_statement&pm=13112&rd=15852

問題概要

　N 人の従業員がいる会社がある。0 番の従業員以外の各従業員には直属の上司が丁度一人いて、i 番の従業員の上司は superior[ i ] である。
　また、各従業員は丁度一つの部署を受け持っている。ある従業員 x の部署に所属する従業員は superior[ y ] = x となる従業員 y である（推移律は成り立たない）。
　更に、各従業員はそれぞれ一つの種類の仕事を担当していて、i 番の従業員の担当する仕事の種類は workType[ i ] である。ある部署について、そこに所属する全ての授業員の仕事の種類が相異なるとき、その部署は好ましい部署であるとする。好ましい部署の個数を求めよ。

解法

　各部署毎に、その部署が好ましいそれか否かチェックして数え上げることで解くことができます。より詳細には、

• 各部署に属する従業員を列挙
• 仕事の種類が相異なるか否かチェック

という流れになります。
　部署毎の従業員の列挙は、各部署毎に列挙用のリスト（配列）を用意して、

• 全ての i について、i 番のリストに i を追加
• 1 以上の i について superior[ i ] のリストに i を追加

すれば列挙できます。
　ある部署内の従業員の仕事の種類が相異なるか否かは、該当する従業員の workType を全て set に入れてから、その size が部署に所属する従業員の数と等しいか否かを調べれば分かります。

コード

using VI = vector<int>;
using VVI = vector<VI>;

#define REP( i, m, n ) for ( int i = (int)( m ); i < (int)( n ); ++i )
#define FOR( e, c ) for ( auto &e : c )

#define PB( n ) push_back( n )

class GoodCompanyDivTwo
{
public:
	int countGood( vector <int> superior, vector <int> workType )
	{
		const int N = superior.size();

		VVI G( N );
		REP( v, 1, N )
		{
			G[ superior[v] ].PB( v );
		}

		int res = 0;
		REP( v, 0, N )
		{
			const int M = G[v].size() + 1;
			G[v].PB( v );

			set<int> works;
			FOR( u, G[v] )
			{
				works.insert( workType[u] );
			}

			res += M == works.size();
		}

		return res;
	}
};