Problem p02671 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>

using namespace std;

const int N = 6e5 + 5;

char s[N];
int p[N], from[N], last[N];
int dig[N];

int main() {
  int n;
  scanf("%d", &n);
  scanf("%s", s);
  int len = strlen(s);
  int pw = 1;
  fill(last, last + len, 0);
  for (int i = 0; i < n; ++i) {
    bool swapped = false;
    for (int j = 0; j < pw; ++j) {
      dig[j] = 0;
      dig[pw + j] = 1;
      dig[pw * 2 + j] = 2;
    }
    for (int j = 0; j < len; ++j) {
      if (s[j] == 'S') {
        swapped = !swapped;
      } else {
        int add = 0;
        for (int k = 0; k < 3; ++k) {
          int now = k * pw + last[j];
          if (swapped) {
            if (dig[now] == 1)
              add = k;
            if (dig[now] == 0)
              dig[now] = 2;
            else if (dig[now] == 1)
              dig[now] = 0;
            else
              dig[now] = 1;
          } else {
            if (dig[now] == 2)
              add = k;
            dig[now] = (dig[now] + 1) % 3;
          }
        }
        last[j] += add * pw;
      }
    }
    for (int j = 0; j < pw; ++j) {
      int now = p[j];
      for (int k = 0; k < 3; ++k) {
        int cur = j + k * pw;
        if (swapped && dig[cur] != 0) {
          dig[cur] = 3 - dig[cur];
        }
        p[cur] = now + dig[cur] * pw;
      }
    }
    pw *= 3;
  }
  for (int i = 0; i < pw; ++i) {
    printf("%d%c", p[i], i + 1 == pw ? '\\n' : ' ');
  }
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

const int N = 6e5 + 5;

char s[N];
int p[N], from[N], last[N];

int main() {
  int n;
  scanf("%d", &n);
  scanf("%s", s);
  int len = strlen(s);
  int pw = 1;
  fill(last, last + len, 0);
  for (int i = 0; i < n; ++i) {
    set<int> st[3];
    for (int j = 0; j < pw; ++j) {
      st[0].insert(j);
      st[1].insert(pw + j);
      st[2].insert(pw * 2 + j);
    }
    for (int j = 0; j < len; ++j) {
      // cerr << i << " " << j << endl;
      if (s[j] == 'S') {
        st[1].swap(st[2]);
      } else {
        int add = 0;
        for (int k = 0; k < 3; ++k) {
          int now = k * pw + last[j];
          for (int z = 0; z < 3; ++z) {
            auto it = st[z].find(now);
            if (it == st[z].end())
              continue;
            if (z == 2)
              add = k;
            st[z].erase(it);
            int y = z + 1;
            if (y == 3)
              y = 0;
            st[y].insert(now);
            break;
          }
        }
        last[j] += add * pw;
      }
      // for (int k = 0; k < 3; ++k) {
      //   cerr << " isi " << k << " : ";
      //   for (int x : st[k])
      //     cerr << x << " ";
      //   cerr << endl;
      // }
    }
    for (int j = 0; j < pw; ++j) {
      int now = p[j];
      // cerr << j << " p[j] " << now << endl;
      for (int k = 0; k < 3; ++k) {
        int cur = j + k * pw;
        for (int z = 0; z < 3; ++z) {
          if (st[z].count(cur)) {
            p[cur] = now + z * pw;
            // cerr << " assign " << cur << " " << k << " " << z << " " <<
            // p[cur] << endl;
            break;
          }
        }
      }
    }
    pw *= 3;
    // for (int j = 0; j < pw; ++j)
    //   cerr << p[j] << " | ";
    // cerr << endl;
  }
  for (int i = 0; i < pw; ++i) {
    printf("%d%c", p[i], i + 1 == pw ? '\\n' : ' ');
  }
  return 0;
}