Problem p02749 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
#define fs first
#define sc second
#define pb push_back
#define mp make_pair
#define eb emplace_back
#define ALL(A) A.begin(), A.end()
#define RALL(A) A.rbegin(), A.rend()
typedef long long ll;
typedef pair<ll, ll> P;
const ll mod = 1000000007;
const ll LINF = 1LL << 60;
const int INF = 1 << 30;
int n;
vector<int> v(100001, 0);
vector<vector<int>> g(100001);

void dfs(int x, int parent, int dist) {
  for (int i = 0; i < g[x].size(); i++) {
    if (g[x][i] != parent) {
      v[g[x][i]] = (dist % 2);
      dfs(g[x][i], x, dist + 1);
    }
  }
}

int main() {
  cin >> n;
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    g[a].pb(b);
    g[b].pb(a);
  }
  dfs(1, -1, 1);
  int cnt_red = 0;
  int cnt_blue = 0;
  for (int i = 1; i <= n; i++) {
    // cout << v[i] << " ";
    if (v[i] == 0) {
      cnt_red++;
    } else {
      cnt_blue++;
    }
  }
  vector<int> ans(n + 1, 0);
  vector<int> s0;
  vector<int> s1;
  vector<int> s2;
  for (int i = 1; i <= n; i++) {
    if (i % 3 == 0) {
      s0.pb(i);
    } else if (i % 3 == 1) {
      s1.pb(i);
    } else {
      s2.pb(i);
    }
  }
  if (cnt_red <= n / 3) {
    int p = 0;
    int q = 0;
    int r = 0;
    for (int i = 1; i <= n; i++) {
      if (v[i] == 0) {
        ans[i] = s0[p];
        p++;
      }
    }
    for (int i = 1; i <= n; i++) {
      if (ans[i] == 0) {
        if (p < s0.size()) {
          ans[i] = s0[p];
          p++;
        } else {
          if (q < s1.size()) {
            ans[i] = s1[q];
            q++;

          } else if (r < s2.size()) {
            ans[i] = s2[r];
            r++;
          }
        }
      }
    }
  } else if (cnt_blue <= n / 3) {
    int p = 0;
    int q = 0;
    int r = 0;
    for (int i = 1; i <= n; i++) {
      if (v[i] == 1) {
        ans[i] = s0[p];
        p++;
      }
    }
    for (int i = 1; i <= n; i++) {
      if (ans[i] == 0) {
        if (p < s0.size()) {
          ans[i] = s0[p];
          p++;
        } else {
          if (q < s1.size()) {
            ans[i] = s1[q];
            q++;

          } else if (r < s2.size()) {
            ans[i] = s2[r];
            r++;
          }
        }
      }
    }
  } else {
    int p = 0;
    int q = 0;
    int r = 0;
    for (int i = 1; i <= n; i++) {
      if (v[i] == 0) {
        if (p < s1.size()) {
          ans[i] = s1[p];
          p++;

        } else {
          ans[i] = s0[q];
          q++;
        }
      }
    }
    for (int i = 1; i <= n; i++) {
      if (ans[i] == 0) {
        if (q < s0.size()) {
          ans[i] = s0[q];
          q++;

        } else if (r < s2.size()) {
          ans[i] = s2[r];
          r++;
        }
      }
    }
  }
  for (int i = 1; i <= n; i++) {
    cout << ans[i] << " ";
  }
  cout << endl;

  return 0;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define fs first
#define sc second
#define pb push_back
#define mp make_pair
#define eb emplace_back
#define ALL(A) A.begin(), A.end()
#define RALL(A) A.rbegin(), A.rend()
typedef long long ll;
typedef pair<ll, ll> P;
const ll mod = 1000000007;
const ll LINF = 1LL << 60;
const int INF = 1 << 30;
int n;
vector<int> v(200001, 0);
vector<vector<int>> g(200001);

void dfs(int x, int parent, int dist) {
  for (int i = 0; i < g[x].size(); i++) {
    if (g[x][i] != parent) {
      v[g[x][i]] = (dist % 2);
      dfs(g[x][i], x, dist + 1);
    }
  }
}

int main() {
  cin >> n;
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    g[a].pb(b);
    g[b].pb(a);
  }
  dfs(1, -1, 1);
  int cnt_red = 0;
  int cnt_blue = 0;
  for (int i = 1; i <= n; i++) {
    // cout << v[i] << " ";
    if (v[i] == 0) {
      cnt_red++;
    } else {
      cnt_blue++;
    }
  }
  vector<int> ans(n + 1, 0);
  vector<int> s0;
  vector<int> s1;
  vector<int> s2;
  for (int i = 1; i <= n; i++) {
    if (i % 3 == 0) {
      s0.pb(i);
    } else if (i % 3 == 1) {
      s1.pb(i);
    } else {
      s2.pb(i);
    }
  }
  if (cnt_red <= n / 3) {
    int p = 0;
    int q = 0;
    int r = 0;
    for (int i = 1; i <= n; i++) {
      if (v[i] == 0) {
        ans[i] = s0[p];
        p++;
      }
    }
    for (int i = 1; i <= n; i++) {
      if (ans[i] == 0) {
        if (p < s0.size()) {
          ans[i] = s0[p];
          p++;
        } else {
          if (q < s1.size()) {
            ans[i] = s1[q];
            q++;

          } else if (r < s2.size()) {
            ans[i] = s2[r];
            r++;
          }
        }
      }
    }
  } else if (cnt_blue <= n / 3) {
    int p = 0;
    int q = 0;
    int r = 0;
    for (int i = 1; i <= n; i++) {
      if (v[i] == 1) {
        ans[i] = s0[p];
        p++;
      }
    }
    for (int i = 1; i <= n; i++) {
      if (ans[i] == 0) {
        if (p < s0.size()) {
          ans[i] = s0[p];
          p++;
        } else {
          if (q < s1.size()) {
            ans[i] = s1[q];
            q++;

          } else if (r < s2.size()) {
            ans[i] = s2[r];
            r++;
          }
        }
      }
    }
  } else {
    int p = 0;
    int q = 0;
    int r = 0;
    for (int i = 1; i <= n; i++) {
      if (v[i] == 0) {
        if (p < s1.size()) {
          ans[i] = s1[p];
          p++;

        } else {
          ans[i] = s0[q];
          q++;
        }
      }
    }
    for (int i = 1; i <= n; i++) {
      if (ans[i] == 0) {
        if (q < s0.size()) {
          ans[i] = s0[q];
          q++;

        } else if (r < s2.size()) {
          ans[i] = s2[r];
          r++;
        }
      }
    }
  }
  for (int i = 1; i <= n; i++) {
    cout << ans[i] << " ";
  }
  cout << endl;

  return 0;
}