Problem p02749 (C++)

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

#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(a) (a).begin(), (a).end()

typedef long long lint;

using namespace std;

struct generator {
  int n;
  int i[3] = {1, 0, 0};
  generator(int _n) : n(_n) {}
  bool remain(int a) {
    if (a > 2)
      return false;
    if (3 * i[a] + a > n)
      return false;
    return true;
  }
  int get(int a) {
    if (!remain(a))
      return 0;
    int ret = 3 * i[a] + a;
    i[a]++;
    return ret;
  }
};

int main() {
  int N;
  cin >> N;
  vector<vector<int>> c(N);
  REP(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    c[a].push_back(b);
    c[b].push_back(a);
  }
  vector<int> d(N, -1);
  int no = 1, ne = 0;
  queue<int> q;
  q.push(0);
  d[0] = 0;
  while (!q.empty()) {
    int p = q.front();
    q.pop();
    for (auto e : c[p])
      if (d[e] < 0) {
        q.push(e);
        d[e] = d[p] + 1;
        if (d[e] % 2 == 0)
          ne++;
        else
          no++;
      }
  }
  REP(i, N) cerr << d[i] << " ";
  cerr << endl;
  cerr << ne << " " << no << endl;
  auto g = generator(N);
  vector<int> ans(N);
  REP(i, N) {
    if (ne <= N / 3) {
      if (d[i] % 2 == 0) {
        ans[i] = g.get(0);
        if (ans[i] == 0)
          return 1;
      } else {
        if (g.remain(1))
          ans[i] = g.get(1);
        else if (g.remain(2))
          ans[i] = g.get(2);
        else
          ans[i] = g.get(0);
      }
    } else if (no <= N / 3) {
      if (d[i] % 2 == 0) {
        if (g.remain(1))
          ans[i] = g.get(1);
        else if (g.remain(2))
          ans[i] = g.get(2);
        else
          ans[i] = g.get(0);
      } else {
        ans[i] = g.get(0);
        if (ans[i] == 0)
          return 1;
      }
    } else {
      if (d[i] % 2 == 0) {
        if (g.remain(1))
          ans[i] = g.get(1);
        else
          ans[i] = g.get(0);
      } else {
        if (g.remain(2))
          ans[i] = g.get(2);
        else
          ans[i] = g.get(0);
      }
    }
  }
  REP(i, N) cout << ans[i] << " ";
  cout << endl;
  return 0;
}


Fixed submission:
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string.h>
#include <vector>

#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(a) (a).begin(), (a).end()

typedef long long lint;

using namespace std;

struct generator {
  int n;
  int i[3] = {1, 0, 0};
  generator(int _n) : n(_n) {}
  bool remain(int a) {
    if (a > 2)
      return false;
    if (3 * i[a] + a > n)
      return false;
    return true;
  }
  int get(int a) {
    if (!remain(a))
      return 0;
    int ret = 3 * i[a] + a;
    i[a]++;
    return ret;
  }
};

int main() {
  int N;
  cin >> N;
  vector<vector<int>> c(N);
  REP(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    c[a].push_back(b);
    c[b].push_back(a);
  }
  vector<int> d(N, -1);
  int no = 0, ne = 1;
  queue<int> q;
  q.push(0);
  d[0] = 0;
  while (!q.empty()) {
    int p = q.front();
    q.pop();
    for (auto e : c[p])
      if (d[e] < 0) {
        q.push(e);
        d[e] = d[p] + 1;
        if (d[e] % 2 == 0)
          ne++;
        else
          no++;
      }
  }
  REP(i, N) cerr << d[i] << " ";
  cerr << endl;
  cerr << ne << " " << no << endl;
  auto g = generator(N);
  vector<int> ans(N);
  REP(i, N) {
    if (ne <= N / 3) {
      if (d[i] % 2 == 0) {
        ans[i] = g.get(0);
        if (ans[i] == 0)
          return 1;
      } else {
        if (g.remain(1))
          ans[i] = g.get(1);
        else if (g.remain(2))
          ans[i] = g.get(2);
        else
          ans[i] = g.get(0);
      }
    } else if (no <= N / 3) {
      if (d[i] % 2 == 0) {
        if (g.remain(1))
          ans[i] = g.get(1);
        else if (g.remain(2))
          ans[i] = g.get(2);
        else
          ans[i] = g.get(0);
      } else {
        ans[i] = g.get(0);
        if (ans[i] == 0)
          return 1;
      }
    } else {
      if (d[i] % 2 == 0) {
        if (g.remain(1))
          ans[i] = g.get(1);
        else
          ans[i] = g.get(0);
      } else {
        if (g.remain(2))
          ans[i] = g.get(2);
        else
          ans[i] = g.get(0);
      }
    }
  }
  REP(i, N) cout << ans[i] << " ";
  cout << endl;
  return 0;
}