okayunのプログラミング日記

主に自分の解いた問題やCodinGameなどの解法などを載せていきます。

ABC082-D : FT Robot

ABC082-D : FT Robot

問題の概要はこちら

最初は O(N3) の dp(dp[i][x][y] := i 番目の命令の後に位置 (x, y) にいるかどうか)だと思ったんですが, これでは MLE かつ TLE してしまいます(工夫すれば MLE は避けられますが計算量はそのまま).
しかし, 'F' は現在向いている方向に 1 つ進むだけなので, 連続する 'F' はまとめることが出来ます. このとき 'T' は区切り文字として見てやります.


例:

"FFFTFFTTFFF" -> "FFF, FF,  , FFF"-> [3, 2, 0, 3] (各数字は 'T' で区切られた区間における 'F' の個数(移動距離)を表す)

こうして出来た数列を S とします. さらに考察すると, S の偶数番目は x 軸方向への移動, 奇数番目は y 軸方向への移動を表すことがわかります. つまり, x 軸方向への移動と y 軸方向への移動は別々に考えることが出来ます. これを利用すると, x 軸方向への移動に関して次のような dp を考えることが出来ます.

dp[i][x] := i 回目の x 軸方向への移動の後, 位置 x にいるかどうか

y 軸方向への移動に関しても同様です. ただし, x 軸方向への移動に関しては, 最初の移動だけ必ず方向が決まっているので実装の際は注意が必要です.

とまぁ, 正しく説明できたのかはわかりませんがこんな感じで解けました(解説見たり某氏からアドバイスもらったりした).
以下が AC したソースコードです.

#include <iostream>
#include <vector>
#include <algorithm>
#include <map>
#include <unordered_map>
#include <functional>
#include <ostream>
#include <queue>
#include <random>
#include <string>
#include <stack>
#include <cassert>
#include <cstring>
#include <cstdlib>
#include <ctime>
#include <cmath>

using std::cin;
using std::cout;
using std::cerr;
using std::endl;

using ll = long long;

template<typename T>
std::ostream& operator << (std::ostream& os, std::vector<T> v) {
  os << "[";
  for (int i = 0; i < int(v.size()); ++i) {
    os << v[i] << (i == int(v.size()) - 1 ? "" : ", ");
  }
  cerr << "]";
  return os;
}

int main() {
  cin.tie(0);
  std::ios::sync_with_stdio(false);

  std::string s;
  int x, y;

  static bool dp[4010][16010]; // dp[i][d] := i回目の移動でdに到達できるか
  std::vector<int> xd, yd;

  cin >> s;
  cin >> x >> y;

  bool ans = true;
  int cnt = 0;

  // s を圧縮
  for (int i = 0; i < int(s.size()); ++i) {
    if (s[i] == 'T') {
      if (xd.size() == yd.size()) {
        xd.push_back(cnt);
      }
      else {
        yd.push_back(cnt);
      }
      cnt = 0;
    }
    else {
      cnt++;

      if (int(s.size()) - 1 == i) {
        if (xd.size() == yd.size()) {
          xd.push_back(cnt);
        }
        else {
          yd.push_back(cnt);
        }
      }
    }
  }

  std::memset(dp, false, sizeof dp);

  // x 座標に関して, s を実行した後に目的の座標に到達できるか
  dp[0][8000] = dp[1][xd[0] + 8000] = true;

  for (int i = 1; i < int(xd.size()); ++i) {
    for (int _x = -8000; _x <= 8000; ++_x) {
      if (dp[i][_x + 8000]) {
        int nx1 = _x + xd[i] + 8000, nx2 = _x - xd[i] + 8000;
        dp[i + 1][nx1] = dp[i + 1][nx2] = true;
      }
    }
  }

  ans &= dp[xd.size()][x + 8000];

  std::memset(dp, false, sizeof dp);

  // y 座標に関しても同様
  dp[0][8000] = true;

  for (int i = 0; i < int(yd.size()); ++i) {
    for (int _y = -8000; _y <= 8000; ++_y) {
      if (dp[i][_y + 8000]) {
        int ny1 = _y + yd[i] + 8000, ny2 = _y - yd[i] + 8000;
        dp[i + 1][ny1] = dp[i + 1][ny2] = true;
      }
    }
  }

  ans &= dp[yd.size()][y + 8000];

  if (ans) {
    cout << "Yes" << endl;
  }
  else {
    cout << "No" << endl;
  }

  return 0;
}

+8000を書かずにスマートに書きたい…
また, 記事に関する指摘などありましたら twitter : @okayunonaha までよろしくお願いします.
(あと Markdown 良い&良い)