👋

三井住友信託銀行プログラミングコンテスト2019 | E - Colorful Hats 2

2020/12/17に公開

競プロ

組み合わせ

色

tech

問題

考えたこと

最初、1番目の人に仮に赤を割り当てた場合、2番目に青を割り当てた場合というふうに考えたが別々の色のパターンで複数の組み合わせができうまくいかないことに気づいた。

色を特定せず各色の帽子の数別の数を管理してその組み合わせをかけてあげればいい。

例えば以下の場合を考える。

$cnt[i]$ は帽子をかぶっている人数 $i$ になりえる色の数である。最初はどの色も0人なので $cnt[0] = 3$ となる。

$A_1$ が0なので、 $A_1$ の人が帽子をかぶると1色で1人帽子をかぶっているパターンが増え、その分0人のパターンが減るので、 $cnt[0] = 2, cnt[1] = 1$ になる。組み合わせは $cnt[0] = 3$ の中から選ぶので3増える。このように順に記録していくと合計の組み合わせ数を求める事ができる。もし $cnt[i] = 0$ の場合は該当する組み合わせがないので答えは0になる。

コード

#include <bits/stdc++.h>

#include <atcoder/all>

using namespace std;
using namespace atcoder;
using ll = long long;
using ld = long double;
using uint = unsigned int;
using ull = unsigned long long;
const int MOD = 1e9 + 7;

using mint = modint1000000007;

int main() {
  ll n;
  cin >> n;
  vector<int> a(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }
  vector<int> cnt(n);  // cnt[i]: 帽子をかぶっている人数iになりえる色の数
  cnt[0] = 3;          // 最初はどの色も0人
  mint ans = 1;
  for (int i = 0; i < n; i++) {
    ll v = cnt[a[i]];
    if (v == 0) {
      cout << 0 << endl;
      return 0;
    }
    ans *= v;
    cnt[a[i]]--;
    cnt[a[i] + 1]++;
  }
  cout << ans.val() << endl;
}

三井住友信託銀行プログラミングコンテスト2019 | E - Colorful Hats 2

問題

考えたこと

コード

参考

Discussion