第二回全国統一プログラミング王決定戦予選 B - 今日から始める競プロ日記
はじめに
日本経済新聞社主催のコンテスト、 第二回全国統一プログラミング王決定戦予選に参加しました。
B問題で無限に時間を使ってしまったので、 同じ過ちを繰り返さないために、 ここにまとめておきます（100分くらいB問題に使った）。
B - Counting of Trees
- D[1] ～ D[N] までの整数列が与えらえる
- 次の条件を満たす、1 ～ Nまでの頂点を持つ木の個数を数える
- 条件：1以上N以下の任意の整数iに対して、頂点1と頂点iの距離がD[i]である
- また、解は998244353の余りを出力
解き方
まず、そもそも木が作れないパターンを判定する。 以下を満たす場合は'0'を出力して終了。
- (1) 頂点1から頂点1までの距離が0かどうか判定
- (2) 距離が0の頂点が二つ以上あってもだめだから判定
- (3) 距離が飛び飛びになっていてもだめだから判定（距離4の頂点はないけど、距離5の頂点はあるみたいな状況）
(3)は判定しなくてもちゃんと実装できていれば自然に回避できるらしい。
次に、各距離ごとの頂点の個数を数える（辺の個数でもある）。
最後に、距離iの頂点の個数をcnt[i]、
距離i+1の頂点の個数をcnt[i+1]として、
cnt[i]のcnt[i+1]乗を掛け合わせていく。
サンプルケース3を例にとる。
- 距離0の頂点が1つ
- 距離1の頂点が2つ
- 距離2の頂点が3つ
- 距離3の頂点が1つ
距離0の頂点から距離1の頂点へ辺を張る方法は必ず1通りしかないのでこれは考えなくてよい。
距離1の2つの頂点から、距離2の3つの頂点への辺の張り方は、2^3 = 8である。
（これは高校数学の組み分け問題の考え方。3つの異なるボールを2人で分けるとき、分け方は何通りかみたいな問題。
参考：【高校数学A】組分け問題全パターン | 受験の月）
同じように考えて、距離2の3つの頂点から、距離3の1つの頂点への辺の張り方は、3^1 = 3通りである。
よって答えは
1 * 8 * 3 = 24
となる。
そして最後に最も重要なことがある。
それはmodを取ること。
私はmodを取ることはちゃんと意識していたけど、 べき乗でpowを使ってそのpowの後でmodを取るようにしていたので、 オーバーフローを抑えられなかったらしい。
コンテスト残り5分くらいのときに、 「あれ、これ掛け算一回やるたびにmod取らないとダメじゃない？」 ということに気が付きました。
modは掛け算を一度行うたびにとったほうが良さそう
コード
#include <bits/stdc++.h>
using namespace std;
int main()
{
// 入力
int N;
cin >> N;
vector<int64_t> D(N); // 与えられる整数列
vector<int64_t> cnt(N, 0); // 各距離の個数をカウント
for(int i = 0; i < N; i++)
{
int t;
cin >> t;
D[i] = t;
cnt[t]++;
}
// 頂点1から頂点1までの距離が0だったらダメ
if(D[0] != 0)
{
cout << "0" << endl;
return 0;
}
// 距離が0である頂点が二つ以上あったらダメ
int zeroCnt = 0;
for(int i = 0; i < D.size(); i++)
{
if(D[i] == 0) zeroCnt++;
if(zeroCnt >= 2)
{
cout << "0" << endl;
return 0;
}
}
// 距離が飛び飛びになっていたらダメ
// （そんな状況はない？）
bool flag = true;
for(int i = 0; i < cnt.size(); i++)
{
if(flag == true && cnt[i] == 0) flag = false;
else if(flag == false && cnt[i] > 0)
{
cout << "0" << endl;
return 0;
}
}
// 頂点1から距離がiの頂点の個数をT[i]
// 頂点1から距離がi+1の頂点の個数をT[i+1]
// ans *= T[i] ^ T[i+1]
// 高校数学の組み分け問題と同じ考え方
// 注意点：一回掛け算するたびにmodをとる
long long ans = 1;
const long long mod = 998244353;
for(int i = 0; i < cnt.size()-1; i++)
{
if(cnt[i] != 0 && cnt[i+1] != 0)
{
for(long long p = 0; p < cnt[i+1]; p++)
{
ans *= cnt[i];
ans %= mod;
}
}
}
// 解を出力
cout << ans << endl;
return 0;
}
おわりに
何回もmodに苦しんでいる。悲しい。
