Problem p00817 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <utility>
#include <vector>

using namespace std;
typedef long long ll;
typedef pair<int, int> PII;
typedef vector<int> VI;
typedef vector<VI> VVI;
#define MP make_pair
#define PB push_back
#define inf 1000000007
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)

int p[302];
int r[302];
int s[302];

void init(int n) {
  for (int i = 0; i < n; ++i) {
    p[i] = i;
    r[i] = 0;
    s[i] = 1;
  }
}

int find(int x) {
  if (x != p[x]) {
    p[x] = find(p[x]);
  }
  return p[x];
}

void union_set(int x, int y) {
  x = find(x);
  y = find(y);
  if (x == y)
    return;
  if (r[x] < r[y]) {
    p[x] = y;
    s[y] += s[x];
  } else {
    p[y] = x;
    s[x] += s[y];
    if (r[x] == r[y])
      r[x]++;
  }
  return;
}

int main() {
  int n, m, p1, p2;
  while (cin >> m >> p1 >> p2) {
    if (m == 0 && p1 == 0 && p2 == 0)
      break;
    if (p1 == 0) {
      for (int i = 0; i < m; i++) {
        int x, y;
        string s;
        cin >> x >> y >> s;
        x--;
        y--;
      }
      cout << "end" << endl;
    } else if (p2 == 0) {
      for (int i = 0; i < m; i++) {
        int x, y;
        string s;
        cin >> x >> y >> s;
        x--;
        y--;
      }
      for (int i = 1; i <= p1; i++) {
        cout << i << endl;
      }
      cout << "end" << endl;
    } else {
      int n = p1 + p2;
      init(n);
      vector<pair<int, int>> node;
      for (int i = 0; i < m; i++) {
        int x, y;
        string s;
        cin >> x >> y >> s;
        x--;
        y--;
        if (s == "no") {
          node.PB(MP(x, y));
        } else {
          union_set(x, y);
        }
      }
      vector<vector<int>> g(n);
      for (auto x : node) {
        g[find(x.first)].PB(find(x.second));
        g[find(x.second)].PB(find(x.first));
      }
      vector<pair<int, int>> pa;
      vector<bool> flag(n);
      for (int i = 0; i < n; i++) {
        if (find(i) != i) {
          flag[i] = 1;
        }
        if (!flag[i]) {
          queue<pair<int, int>> que;
          vector<int> c1, c2;
          que.push(MP(i, 0));
          flag[i] = 1;
          while (!que.empty()) {
            auto x = que.front();
            flag[x.first] = 1;
            que.pop();
            if (x.second % 2 == 0) {
              c1.PB(x.first);
            } else {
              c2.PB(x.first);
            }
            for (auto y : g[x.first]) {
              if (flag[y] == 0) {
                que.push(MP(y, x.second + 1));
              }
            }
          }
          for (int j = 1; j < c1.size(); j++) {
            union_set(c1[0], c1[j]);
          }
          for (int j = 1; j < c2.size(); j++) {
            union_set(c2[0], c2[j]);
          }
          if (c2.size() == 0) {
            pa.PB(MP(find(c1[0]), -1));
          } else {
            pa.PB(MP(find(c1[0]), find(c2[0])));
          }
        }
      }
      // for(auto x:pa){
      // 	cout << x.first << " " << s[x.first] <<" " << x.second << " " <<
      // s[x.second] << endl;
      // }
      int dp[601][601] = {};
      dp[0][0] = 1;
      int k = pa.size();
      for (int i = 0; i < k; i++) {
        auto x = pa[i];
        if (x.second == -1) {
          for (int j = 0; j <= 600; j++) {
            if (dp[i][j] != 0) {
              dp[i + 1][j + s[x.first]] += dp[i][j];
              dp[i + 1][j] += dp[i][j];
            }
          }
        } else {
          for (int j = 0; j <= 600; j++) {
            if (dp[i][j] != 0) {
              dp[i + 1][j + s[x.first]] += dp[i][j];
              dp[i + 1][j + s[x.second]] += dp[i][j];
            }
          }
        }
      }
      if (dp[k][p1] != 1) {
        cout << "no" << endl;
      } else {
        set<int> st;
        int pos = p1;
        for (int i = k - 1; i >= 0; i--) {
          auto x = pa[i];
          if (x.second == -1) {
            if (pos - s[x.first] >= 0) {
              if (dp[i][pos - s[x.first]] != 0) {
                pos = pos - s[x.first];
                st.insert(x.first);
              }
            }
          } else {
            bool fflag = 1;
            if (pos - s[x.first] >= 0) {
              if (dp[i][pos - s[x.first]] != 0) {
                pos = pos - s[x.first];
                st.insert(x.first);
                fflag = 0;
              }
            }
            if (fflag && pos - s[x.second] >= 0) {
              if (dp[i][pos - s[x.second]] != 0) {
                pos = pos - s[x.second];
                st.insert(x.second);
              }
            }
          }
        }
        for (int i = 0; i < n; i++) {
          if (st.count(find(i))) {
            cout << i + 1 << endl;
          }
        }
        cout << "end" << endl;
      }
    }
  }

  return 0;
}


Fixed submission:
#include <algorithm>
#include <cmath>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <utility>
#include <vector>

using namespace std;
typedef long long ll;
typedef pair<int, int> PII;
typedef vector<int> VI;
typedef vector<VI> VVI;
#define MP make_pair
#define PB push_back
#define inf 1000000007
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)

int p[602];
int r[602];
int s[602];

void init(int n) {
  for (int i = 0; i < n; ++i) {
    p[i] = i;
    r[i] = 0;
    s[i] = 1;
  }
}

int find(int x) {
  if (x != p[x]) {
    p[x] = find(p[x]);
  }
  return p[x];
}

void union_set(int x, int y) {
  x = find(x);
  y = find(y);
  if (x == y)
    return;
  if (r[x] < r[y]) {
    p[x] = y;
    s[y] += s[x];
  } else {
    p[y] = x;
    s[x] += s[y];
    if (r[x] == r[y])
      r[x]++;
  }
  return;
}

int main() {
  int n, m, p1, p2;
  while (cin >> m >> p1 >> p2) {
    if (m == 0 && p1 == 0 && p2 == 0)
      break;
    if (p1 == 0) {
      for (int i = 0; i < m; i++) {
        int x, y;
        string s;
        cin >> x >> y >> s;
        x--;
        y--;
      }
      cout << "end" << endl;
    } else if (p2 == 0) {
      for (int i = 0; i < m; i++) {
        int x, y;
        string s;
        cin >> x >> y >> s;
        x--;
        y--;
      }
      for (int i = 1; i <= p1; i++) {
        cout << i << endl;
      }
      cout << "end" << endl;
    } else {
      int n = p1 + p2;
      init(n);
      vector<pair<int, int>> node;
      for (int i = 0; i < m; i++) {
        int x, y;
        string s;
        cin >> x >> y >> s;
        x--;
        y--;
        if (s == "no") {
          node.PB(MP(x, y));
        } else {
          union_set(x, y);
        }
      }
      vector<vector<int>> g(n);
      for (auto x : node) {
        g[find(x.first)].PB(find(x.second));
        g[find(x.second)].PB(find(x.first));
      }
      vector<pair<int, int>> pa;
      vector<bool> flag(n);
      for (int i = 0; i < n; i++) {
        if (find(i) != i) {
          flag[i] = 1;
        }
        if (!flag[i]) {
          queue<pair<int, int>> que;
          vector<int> c1, c2;
          que.push(MP(i, 0));
          flag[i] = 1;
          while (!que.empty()) {
            auto x = que.front();
            flag[x.first] = 1;
            que.pop();
            if (x.second % 2 == 0) {
              c1.PB(x.first);
            } else {
              c2.PB(x.first);
            }
            for (auto y : g[x.first]) {
              if (flag[y] == 0) {
                que.push(MP(y, x.second + 1));
              }
            }
          }
          for (int j = 1; j < c1.size(); j++) {
            union_set(c1[0], c1[j]);
          }
          for (int j = 1; j < c2.size(); j++) {
            union_set(c2[0], c2[j]);
          }
          if (c2.size() == 0) {
            pa.PB(MP(find(c1[0]), -1));
          } else {
            pa.PB(MP(find(c1[0]), find(c2[0])));
          }
        }
      }
      // for(auto x:pa){
      // 	cout << x.first << " " << s[x.first] <<" " << x.second << " " <<
      // s[x.second] << endl;
      // }
      int dp[601][601] = {};
      dp[0][0] = 1;
      int k = pa.size();
      for (int i = 0; i < k; i++) {
        auto x = pa[i];
        if (x.second == -1) {
          for (int j = 0; j <= 600; j++) {
            if (dp[i][j] != 0) {
              dp[i + 1][j + s[x.first]] += dp[i][j];
              dp[i + 1][j] += dp[i][j];
            }
          }
        } else {
          for (int j = 0; j <= 600; j++) {
            if (dp[i][j] != 0) {
              dp[i + 1][j + s[x.first]] += dp[i][j];
              dp[i + 1][j + s[x.second]] += dp[i][j];
            }
          }
        }
      }
      if (dp[k][p1] != 1) {
        cout << "no" << endl;
      } else {
        set<int> st;
        int pos = p1;
        for (int i = k - 1; i >= 0; i--) {
          auto x = pa[i];
          if (x.second == -1) {
            if (pos - s[x.first] >= 0) {
              if (dp[i][pos - s[x.first]] != 0) {
                pos = pos - s[x.first];
                st.insert(x.first);
              }
            }
          } else {
            bool fflag = 1;
            if (pos - s[x.first] >= 0) {
              if (dp[i][pos - s[x.first]] != 0) {
                pos = pos - s[x.first];
                st.insert(x.first);
                fflag = 0;
              }
            }
            if (fflag && pos - s[x.second] >= 0) {
              if (dp[i][pos - s[x.second]] != 0) {
                pos = pos - s[x.second];
                st.insert(x.second);
              }
            }
          }
        }
        for (int i = 0; i < n; i++) {
          if (st.count(find(i))) {
            cout << i + 1 << endl;
          }
        }
        cout << "end" << endl;
      }
    }
  }

  return 0;
}