Problem p03812 (C++)

Buggy submission (status: Runtime Error):


#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;

#define REP(i, n) for (int i = 1; i <= int(n); i++)
#define MP make_pair
#define PB push_back
#define SZ(x) (int((x).size()))
#define ALL(x) (x).begin(), (x).end()
#define X first
#define Y second
template <typename T> inline bool chkmin(T &a, const T &b) {
  return a > b ? a = b, 1 : 0;
}
template <typename T> inline bool chkmax(T &a, const T &b) {
  return a < b ? a = b, 1 : 0;
}

typedef long long LL;
typedef pair<LL, int> pii;
typedef long double LD;
const int INF = 0x3f3f3f3f;

template <class T> inline bool RD(T &ret) {
  char c;
  int sgn;
  if (c = getchar(), c == EOF)
    return 0;
  while (c != '-' && (c < '0' || c > '9'))
    c = getchar();
  sgn = (c == '-') ? -1 : 1, ret = (c == '-') ? 0 : (c - '0');
  while (c = getchar(), c >= '0' && c <= '9')
    ret = ret * 10 + (c - '0');
  ret *= sgn;
  return 1;
}
template <class T> inline void PT(T x) {
  if (x < 0)
    putchar('-'), x = -x;
  if (x > 9)
    PT(x / 10);
  putchar(x % 10 + '0');
}

const int N = 1e3 + 100;
set<int> from[N];
set<int> to[N];
int val[N];
queue<int> que;
int deg[N], ans[N];

void updatewin(int u) {
  for (auto v : from[u]) {
    if (ans[v] != -1)
      continue;
    if (val[v] > val[u]) {
      deg[v]--;
      from[v].insert(u);
      if (deg[v] == 0)
        que.push(v);
    }
  }
}

void updatefail(int u) {
  for (auto v : from[u]) {
    if (ans[v] != -1)
      continue;
    if (val[v] > val[u]) {
      ans[v] = 1;
      updatewin(v);
    }
  }
}
int main() {
  memset(ans, -1, sizeof(ans));
  int n;
  cin >> n;
  for (int i = 1; i <= n; i++)
    scanf("%d", &val[i]);
  for (int i = 1; i < n; i++) {
    int u, v;
    scanf("%d %d", &u, &v);
    if (val[u] <= val[v]) {
      from[u].insert(v);
      to[v].insert(u);
    }
    if (val[u] >= val[v]) {
      from[v].insert(u);
      to[u].insert(v);
    }
  }
  for (int u = 1; u <= n; u++) {
    int sz = 0;
    for (auto v : to[u]) {
      if (from[u].count(v) == 0)
        deg[u]++;
    }
    if (deg[u] == 0)
      que.push(u);
  }
  while (!que.empty()) {
    int u = que.front();
    que.pop();
    ans[u] = 0;
    updatefail(u);
  }
  vector<int> res;
  for (int i = 1; i <= n; i++)
    if (ans[i] == 1)
      res.push_back(i);
  for (int i = 0; i < SZ(res); i++)
    printf("%d%c", res[i], i == SZ(res) - 1 ? '\n' : ' ');
}


Fixed submission:


#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;

#define REP(i, n) for (int i = 1; i <= int(n); i++)
#define MP make_pair
#define PB push_back
#define SZ(x) (int((x).size()))
#define ALL(x) (x).begin(), (x).end()
#define X first
#define Y second
template <typename T> inline bool chkmin(T &a, const T &b) {
  return a > b ? a = b, 1 : 0;
}
template <typename T> inline bool chkmax(T &a, const T &b) {
  return a < b ? a = b, 1 : 0;
}

typedef long long LL;
typedef pair<LL, int> pii;
typedef long double LD;
const int INF = 0x3f3f3f3f;

template <class T> inline bool RD(T &ret) {
  char c;
  int sgn;
  if (c = getchar(), c == EOF)
    return 0;
  while (c != '-' && (c < '0' || c > '9'))
    c = getchar();
  sgn = (c == '-') ? -1 : 1, ret = (c == '-') ? 0 : (c - '0');
  while (c = getchar(), c >= '0' && c <= '9')
    ret = ret * 10 + (c - '0');
  ret *= sgn;
  return 1;
}
template <class T> inline void PT(T x) {
  if (x < 0)
    putchar('-'), x = -x;
  if (x > 9)
    PT(x / 10);
  putchar(x % 10 + '0');
}

const int N = 3e3 + 100;
set<int> from[N];
set<int> to[N];
int val[N];
queue<int> que;
int deg[N], ans[N];

void updatewin(int u) {
  for (auto v : from[u]) {
    if (ans[v] != -1)
      continue;
    if (val[v] > val[u]) {
      deg[v]--;
      from[v].insert(u);
      if (deg[v] == 0)
        que.push(v);
    }
  }
}

void updatefail(int u) {
  for (auto v : from[u]) {
    if (ans[v] != -1)
      continue;
    if (val[v] > val[u]) {
      ans[v] = 1;
      updatewin(v);
    }
  }
}
int main() {
  memset(ans, -1, sizeof(ans));
  int n;
  cin >> n;
  for (int i = 1; i <= n; i++)
    scanf("%d", &val[i]);
  for (int i = 1; i < n; i++) {
    int u, v;
    scanf("%d %d", &u, &v);
    if (val[u] <= val[v]) {
      from[u].insert(v);
      to[v].insert(u);
    }
    if (val[u] >= val[v]) {
      from[v].insert(u);
      to[u].insert(v);
    }
  }
  for (int u = 1; u <= n; u++) {
    int sz = 0;
    for (auto v : to[u]) {
      if (from[u].count(v) == 0)
        deg[u]++;
    }
    if (deg[u] == 0)
      que.push(u);
  }
  while (!que.empty()) {
    int u = que.front();
    que.pop();
    ans[u] = 0;
    updatefail(u);
  }
  vector<int> res;
  for (int i = 1; i <= n; i++)
    if (ans[i] == 1)
      res.push_back(i);
  for (int i = 0; i < SZ(res); i++)
    printf("%d%c", res[i], i == SZ(res) - 1 ? '\n' : ' ');
}