TopCoderのDPの問題を解いた

Topcoder SRM 774 Div1 Medium ClassRankings

上記の問題を解こうとしたけど解けなかった。
Topcoderの検索ページでカテゴリ「Dynamic Programming」で検索した際に一番上に表示された問題がこの問題だった。

DPを次のように定義する。

dp[x][a][b][c] := 点数の高い順に生徒の点数を決めていったときの、最後の生徒の点数をx、学校1, 2, 3の残り生徒数をa, b, cとしたときの数え上げの数

遷移は、x - 1がlo[i]以上であればその学校の生徒を選べるので選び、そうでない場合は選ばないようにする。
このとき、学校iの残り生徒数が0のときは選ばないよう注意する。
初期値は、dp[x][0][0][0] = 1。

…という方針で実装したんだけどWAだった。
もしかするとTLEかも。Topcoderの仕様はわからないから何とも言えない。
「Failed」と表示される。けれど具体的な原因はわからない。
手元の環境で50 50 50 10 20 30 160 170 180を入力として与えると10秒以上かかったので、おそらくTLEだと思う。

コードは残しておく。

#include <bits/stdc++.h>
using namespace std;

const int mod = 1000000007;

vector<int> loo, hii;
#define MAP(T) unordered_map<int, T>
MAP(MAP(MAP(MAP(int)))) mp;

int rec(int i, int j, int k, int l) {
  if (j == 0 && k == 0 && l == 0)
    return 1;
  if (mp.count(i) && mp[i].count(j) && mp[i][j].count(k) && mp[i][j][k].count(l)) {
    return mp[i][j][k][l];
  }
  int ret = 0;
  i--;
  if (loo[0] <= i && j > 0) ret += rec(min(i, hii[0]), j - 1, k, l), ret %= mod;
  if (loo[1] <= i && k > 0) ret += rec(min(i, hii[1]), j, k - 1, l), ret %= mod;
  if (loo[2] <= i && l > 0) ret += rec(min(i, hii[2]), j, k, l - 1), ret %= mod;
  i++;
  return mp[i][j][k][l] = ret;
};

struct ClassRankings {
  int countWays(vector<int> amt, vector<int> lo, vector<int> hi) {
    loo = lo;
    hii = hi;
    return rec(3001, amt[0], amt[1], amt[2]);
  }
};

メモ化するデータ構造としてunordered_mapを使っている。

一度通常の配列にしたりもしたが、int dp[3002][51][51][51]だと「Message: caught signal SIGKILL」と実行時エラーになる。
今計算するとメモリ使用量が1.6GBになったから当然か（制約は256MB）。

座標圧縮しようともしたけど、うまくできなかった。
lo, hiの要素の集合から2つを選び、その2つの差が150を超えるとき、差を150まで圧縮することにより状態数は減らせるけど、それでも3000→750に減るだけ。400MBくらい使うので、MLEの制限は抜けられない。

ここで諦めた。いちおう座標圧縮後のコードも残しておく。
main()もあるのでそのままの提出はできない。

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < int(n); i++)
#define reps(i, n) for (int i = 1; i <= int(n); i++)
using namespace std;

const int mod = 1000000007;

vector<int> loo, hii;
int dp[752][51][51][51];

int rec(int i, int j, int k, int l) {
  if (j == 0 && k == 0 && l == 0)
    return 1;
  if (dp[i][j][k][l] != -1)
    return dp[i][j][k][l];
  int ret = 0;
  i--;
  if (loo[0] <= i && j > 0) ret += rec(min(i, hii[0]), j - 1, k, l), ret %= mod;
  if (loo[1] <= i && k > 0) ret += rec(min(i, hii[1]), j, k - 1, l), ret %= mod;
  if (loo[2] <= i && l > 0) ret += rec(min(i, hii[2]), j, k, l - 1), ret %= mod;
  i++;
  return dp[i][j][k][l] = ret;
};

struct ClassRankings {
  int countWays(vector<int> amt, vector<int> lo, vector<int> hi) {
    {
      map<int, int> ma;
      rep(i, 3) ma[lo[i]] = lo[i];
      rep(i, 3) ma[hi[i]] = hi[i];
      vector<int> ve;
      for (auto m : ma) ve.push_back(m.first);
      reps(i, ve.size() - 1) {
        if (ve[i] - ve[i - 1] > 150) {
          ma[ve[i]] = ma[ve[i - 1]] + 150;
        }
      }
      rep(i, 3) lo[i] = ma[lo[i]];
      rep(i, 3) hi[i] = ma[hi[i]];
      int mi = *min_element(lo.begin(), lo.end());
      rep(i, 3) lo[i] -= mi;
      rep(i, 3) hi[i] -= mi;
    }
    rep(i, 752) rep(j, 51) rep(k, 51) rep(l, 51) {
      dp[i][j][k][l] = -1;
    }
    loo = lo;
    hii = hi;
    return rec(751, amt[0], amt[1], amt[2]);
  }
};

int main() {
  vector<int> amt(3), lo(3), hi(3);
  rep(i, 3) cin >> amt[i];
  rep(i, 3) cin >> lo[i];
  rep(i, 3) cin >> hi[i];
  cout << ClassRankings().countWays(amt, lo, hi) << endl;
  return 0;
}