第4回 ドワンゴからの挑戦状 予選 C - Kill/Death

問題ページ C: Kill/Death - 第4回 ドワンゴからの挑戦状 予選 | AtCoder

考えたこと

• sample1を読むと5デスを4人に割り振るようなのを求められればよさそうでこれは分割数なのでできる
• sample2を読むとキル数が違う人ごとにまとめてそれぞれ計算していくとよさそう
• dp[i][j] = (i番目のキル数の人まででjデスを振り分けた場合) みたいなDPが生える
• 状態数O(Nsum(kill))で遷移O(N)くらいで書けそうなので書くと通った

解法

dp[i][j] = (i番目のキル数の人まででjデスを振り分けた場合) としてDPをする。dp[i][j] = sum(dp[i-1][k] * (j-kデスをi番目のキル数の人数に割り振る)) とするとO(N^2 sum(kill)) で10^7くらいなので通る。j-kデスをi番目のキル数の人数に割り振るパターン数は分割数のDPをしておくと前計算できる。

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;
#define int ll
typedef vector<int> VI;
typedef vector<VI> VVI;
typedef pair<int, int> PII;

#define FOR(i, a, n) for (ll i = (ll)a; i < (ll)n; ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(x) x.begin(), x.end()
#define IN(a, b, x) (a<=x&&x<b)
#define PB push_back

const ll LLINF = (1LL<<60);
const int INF = (1LL<<30);
const int MOD = 1000000007;

template <typename T> T &chmin(T &a, const T &b) { return a = min(a, b); }
template <typename T> T &chmax(T &a, const T &b) { return a = max(a, b); }
template<class S,class T>
ostream &operator <<(ostream& out,const pair<S,T>& a){
  out<<'('<<a.first<<','<<a.second<<')';
  return out;
}

int dx[] = {0, 1, 0, -1}, dy[] = {1, 0, -1, 0};

int a[105], b[105];
int dp[1005][1005], dp1[105][1005], dp2[105][1005];
signed main(void)
{
  int n, m;
  cin >> n >> m;
  REP(i, n) cin >> a[i];
  REP(i, m) cin >> b[i];

  VI v1, v2;
  int cnt = 1, sum1 = a[0];
  FOR(i, 1, n) {
    sum1 += a[i];
    if(a[i] == a[i-1]) cnt++;
    else {
      v1.PB(cnt);
      cnt = 1;
    }
  }
  v1.PB(cnt);

  cnt = 1;
  int sum2 = b[0];
  FOR(i, 1, m) {
    sum2 += b[i];
    if(b[i] == b[i-1]) cnt++;
    else {
      v2.PB(cnt);
      cnt = 1;
    }
  }
  v2.PB(cnt);

  // 分割数
  dp[0][0] = 1;
  FOR(i, 1, 1005) {
    FOR(j, 0, 1005) {
      if(j-i >= 0) {
        dp[i][j] = (dp[i-1][j] + dp[i][j-i]) % MOD;
      } else {
        dp[i][j] = dp[i-1][j];
      }
    }
  }

  // v1[0]について
  dp1[0][0] = 1;
  FOR(i, 1, sum2+1) {
    dp1[0][i] = dp[v1[0]][i];
  }

  FOR(i, 1, v1.size()) REP(j, sum2+1) {
    // dp[i-1][k]→dp[i][j]
    REP(k, j+1) {
      // dp[i-1][k] * (j-kをv1[i]人に振り分けたとき)
      (dp1[i][j] += dp1[i-1][k] * dp[v1[i]][j-k] % MOD) %= MOD;
    }
  }

  dp2[0][0] = 1;
  FOR(i, 1, sum1+1) {
    dp2[0][i] = dp[v2[0]][i];
  }
  FOR(i, 1, v2.size()) REP(j, sum1+1) {
    // dp[i-1][k]→dp[i][j]
    REP(k, j+1) {
      // dp[i-1][k] * (j-kをv1[i]人に振り分けたとき)
      (dp2[i][j] += dp2[i-1][k] * dp[v2[i]][j-k] % MOD) %= MOD;
    }
  }

  cout << dp1[v1.size()-1][sum2] * dp2[v2.size()-1][sum1] % MOD << endl;

  return 0;
}