Problem p00817 (C++) — buggy submission, verdict: Memory Limit Exceeded

#include <bits/stdc++.h>

using namespace std;

#define int short

using vint = vector<int>;

struct edge {
  int to, kind;
  edge() {}
  edge(int to, int kind) : to(to), kind(kind) {}
};

using Graph = vector<vector<edge>>;

struct UF {
  vint data;
  UF() {}
  void init(int n) {
    data.clear();
    data.resize(n, -1);
  }
  int find(int x) { return data[x] < 0 ? x : data[x] = find(data[x]); }
  int size(int x) { return -data[find(x)]; }
  bool same(int x, int y) { return find(x) == find(y); }
  void unite(int x, int y) {
    x = find(x), y = find(y);
    if (x == y)
      return;
    if (data[x] > data[y])
      swap(x, y);
    data[x] += data[y];
    data[y] = x;
  }
};

int Q, A, B, N;
Graph graph;
Graph rgraph;
UF uf;

vint col[2];
vint divines[2];

void fillcolor(int s, int c) {
  queue<int> que;
  que.push(s);
  col[c][s] = c;
  if (!c)
    ++divines[c][s];
  while (!que.empty()) {
    int u = que.front();
    que.pop();
    for (edge e : graph[u]) {
      if (col[c][e.to] == -1) {
        que.push(e.to);
        col[c][e.to] = col[c][u] ^ e.kind;
        if (!col[c][e.to])
          ++divines[c][s];
      } else if (col[c][e.to] != (col[c][u] ^ e.kind)) {
        divines[c][s] = -1;
        return;
      }
    }
    for (edge e : rgraph[u]) {
      if (col[c][e.to] == -1) {
        que.push(e.to);
        col[c][e.to] = col[c][u] ^ e.kind;
        if (!col[c][e.to])
          ++divines[c][s];
      } else if (col[c][e.to] != (col[c][u] ^ e.kind)) {
        divines[c][s] = -1;
        return;
      }
    }
  }
  // cout << s << " " << c << " " << divines[c][s] << endl;
}

void init() {
  graph.clear();
  graph.resize(N);
  rgraph.clear();
  rgraph.resize(N);
  col[0].clear();
  col[0].resize(N, -1);
  col[1].clear();
  col[1].resize(N, -1);
  divines[0].clear();
  divines[0].resize(N);
  divines[1].clear();
  divines[1].resize(N);
  uf.init(N);
}

using Pi = pair<int, int>;

// int dp[303][303][303];
Pi rev[303][303][303];

signed main() {
  while (cin >> Q >> A >> B, Q || A || B) {
    N = A + B;
    init();
    for (int i = 0; i < Q; ++i) {
      int a, b;
      string s;
      cin >> a >> b >> s;
      --a, --b;
      graph[a].emplace_back(b, s == "no");
      rgraph[b].emplace_back(a, s == "no");
      uf.unite(a, b);
    }
    // map<int, int> mp;
    vint mp(N);
    int sz = 0;
    for (int i = 0; i < N; ++i) {
      if (uf.find(i) != i)
        continue;
      mp[sz++] = i;
      for (int j = 0; j < 2; ++j) {
        fillcolor(i, j);
      }
    }
    for (int i = 0; i < 303; ++i) {
      for (int j = 0; j < 303; ++j) {
        for (int k = 0; k < 303; ++k) {
          // dp[i][j][k] = 0;
          rev[i][j][k] = Pi(-2, -2);
        }
      }
    }
    // dp[0][0][0] = 1;
    rev[0][0][0] = Pi(-1, -1);
    for (int i = 0; i < sz; ++i) {
      for (int j = 0; j <= A; ++j) {
        for (int k = 0; k <= B; ++k) {
          // if(!dp[i][j][k]) continue;
          if (rev[i][j][k].first == -2)
            continue;
          for (int l = 0; l < 2; ++l) {
            if (divines[l][mp[i]] == -1)
              continue;
            int x = divines[l][mp[i]];
            int y = uf.size(mp[i]) - x;
            if (j + x > A || k + y > B)
              continue;
            // if(!dp[i+1][j+x][k+y]) {
            if (rev[i + 1][j + x][k + y].first == -2) {
              // dp[i+1][j+x][k+y] = 1;
              rev[i + 1][j + x][k + y] = Pi(mp[i], l);
            } else {
              // dp[i+1][j+x][k+y] = 2;
              rev[i + 1][j + x][k + y] = Pi(-1, -1);
            }
          }
        }
      }
    }
    // if(dp[sz][A][B] != 1) {
    if (rev[sz][A][B].first < 0) {
      cout << "no" << endl;
      continue;
    } else {
      vint ans;
      int a = A, b = B;
      for (int i = sz; i > 0; --i) {
        Pi p = rev[i][a][b];
        for (int j = 0; j < N; ++j) {
          if (uf.find(j) != p.first)
            continue;
          if (!col[p.second][j])
            ans.push_back(j);
        }
        a -= divines[p.second][p.first];
        b -= (uf.size(p.first) - divines[p.second][p.first]);
      }
      sort(ans.begin(), ans.end());
      for (int id : ans)
        cout << id + 1 << endl;
      cout << "end" << endl;
    }
  }

  return 0;
}


Error: MLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

#define int short

using vint = vector<int>;

struct edge {
  int to, kind;
  edge() {}
  edge(int to, int kind) : to(to), kind(kind) {}
};

using Graph = vector<vector<edge>>;

struct UF {
  vint data;
  UF() {}
  void init(int n) {
    data.clear();
    data.resize(n, -1);
  }
  int find(int x) { return data[x] < 0 ? x : data[x] = find(data[x]); }
  int size(int x) { return -data[find(x)]; }
  bool same(int x, int y) { return find(x) == find(y); }
  void unite(int x, int y) {
    x = find(x), y = find(y);
    if (x == y)
      return;
    if (data[x] > data[y])
      swap(x, y);
    data[x] += data[y];
    data[y] = x;
  }
};

int Q, A, B, N;
Graph graph;
Graph rgraph;
UF uf;

vint col[2];
vint divines[2];

void fillcolor(int s, int c) {
  queue<int> que;
  que.push(s);
  col[c][s] = c;
  if (!c)
    ++divines[c][s];
  while (!que.empty()) {
    int u = que.front();
    que.pop();
    for (edge e : graph[u]) {
      if (col[c][e.to] == -1) {
        que.push(e.to);
        col[c][e.to] = col[c][u] ^ e.kind;
        if (!col[c][e.to])
          ++divines[c][s];
      } else if (col[c][e.to] != (col[c][u] ^ e.kind)) {
        divines[c][s] = -1;
        return;
      }
    }
    for (edge e : rgraph[u]) {
      if (col[c][e.to] == -1) {
        que.push(e.to);
        col[c][e.to] = col[c][u] ^ e.kind;
        if (!col[c][e.to])
          ++divines[c][s];
      } else if (col[c][e.to] != (col[c][u] ^ e.kind)) {
        divines[c][s] = -1;
        return;
      }
    }
  }
  // cout << s << " " << c << " " << divines[c][s] << endl;
}

void init() {
  graph.clear();
  graph.resize(N);
  rgraph.clear();
  rgraph.resize(N);
  col[0].clear();
  col[0].resize(N, -1);
  col[1].clear();
  col[1].resize(N, -1);
  divines[0].clear();
  divines[0].resize(N);
  divines[1].clear();
  divines[1].resize(N);
  uf.init(N);
}

using Pi = pair<int, int>;

int dp[303][303][303];
Pi rev[303][303][303];

signed main() {
  while (cin >> Q >> A >> B, Q || A || B) {
    N = A + B;
    init();
    for (int i = 0; i < Q; ++i) {
      int a, b;
      string s;
      cin >> a >> b >> s;
      --a, --b;
      graph[a].emplace_back(b, s == "no");
      rgraph[b].emplace_back(a, s == "no");
      uf.unite(a, b);
    }
    map<int, int> mp;
    int sz = 0;
    for (int i = 0; i < N; ++i) {
      if (uf.find(i) != i)
        continue;
      mp[sz++] = i;
      for (int j = 0; j < 2; ++j) {
        fillcolor(i, j);
      }
    }
    for (int i = 0; i < 303; ++i) {
      for (int j = 0; j < 303; ++j) {
        for (int k = 0; k < 303; ++k) {
          dp[i][j][k] = 0;
          rev[i][j][k] = Pi(-1, -1);
        }
      }
    }
    dp[0][0][0] = 1;
    for (int i = 0; i < sz; ++i) {
      for (int j = 0; j <= A; ++j) {
        for (int k = 0; k <= B; ++k) {
          if (!dp[i][j][k])
            continue;
          for (int l = 0; l < 2; ++l) {
            if (divines[l][mp[i]] == -1)
              continue;
            int x = divines[l][mp[i]];
            int y = uf.size(mp[i]) - x;
            if (j + x > A || k + y > B)
              continue;
            if (!dp[i + 1][j + x][k + y]) {
              dp[i + 1][j + x][k + y] = 1;
              rev[i + 1][j + x][k + y] = Pi(mp[i], l);
            } else {
              dp[i + 1][j + x][k + y] = 2;
              rev[i + 1][j + x][k + y] = Pi(-1, -1);
            }
          }
        }
      }
    }
    if (dp[sz][A][B] != 1) {
      cout << "no" << endl;
      continue;
    } else {
      vint ans;
      int a = A, b = B;
      for (int i = sz; i > 0; --i) {
        Pi p = rev[i][a][b];
        for (int j = 0; j < N; ++j) {
          if (uf.find(j) != p.first)
            continue;
          if (!col[p.second][j])
            ans.push_back(j);
        }
        a -= divines[p.second][p.first];
        b -= (uf.size(p.first) - divines[p.second][p.first]);
      }
      sort(ans.begin(), ans.end());
      for (int id : ans)
        cout << id + 1 << endl;
      cout << "end" << endl;
    }
  }

  return 0;
}