Problem p01833 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>

using namespace std;

const int SHIFT = 10000;
const int INF = 1 << 29;

int N, A, B, C, M;
int color[10000], small[10005];
vector<int> g[10005];
pair<int, int> v[3][10005];

void solve(int start, int id) {
  queue<int> que;
  v[id][start] = {0, small[start]};
  que.emplace(start);
  while (!que.empty()) {
    int p = que.front();
    que.pop();
    for (auto &to : g[p]) {
      if (v[id][p].first + 1 == v[id][to].first) {
        v[id][to].second = min(v[id][to].second, v[id][p].second);
      } else if (v[id][p].first + 1 < v[id][to].first) {
        v[id][to] = {v[id][p].first + 1, min(small[to], v[id][p].second)};
        que.emplace(to);
      }
    }
  }
}

int main() {
  memset(color, -1, sizeof(color));
  fill_n(small, 10005, INF);

  scanf("%d %d %d %d", &N, &A, &B, &C);
  for (int i = 0; i < N; i++)
    small[i] = i;
  for (int i = 0; i < A + B + C; i++) {
    int K;
    scanf("%d", &K);
    --K;
    if (i < A)
      color[K] = 0;
    else if (i < A + B)
      color[K] = 1;
    else
      color[K] = 2;
    small[color[K] + SHIFT] = min(small[color[K] + SHIFT], K);
  }

  scanf("%d", &M);
  while (M--) {
    int a, b;
    scanf("%d %d", &a, &b);
    --a, --b;
    if (~color[a])
      a = color[a] + SHIFT;
    if (~color[b])
      b = color[b] + SHIFT;
    g[a].emplace_back(b);
    g[b].emplace_back(a);
  }
  for (int i = 0; i < 10005; i++) {
    sort(begin(g[i]), end(g[i]));
    g[i].erase(unique(begin(g[i]), end(g[i])), end(g[i]));
  }

  fill_n(*v, 3 * 10005, make_pair(INF, INF));
  solve(SHIFT, 0);
  solve(SHIFT + 1, 1);
  solve(SHIFT + 2, 2);

  vector<int> vertex;
  vertex.push_back(SHIFT);
  vertex.push_back(SHIFT + 1);
  vertex.push_back(SHIFT + 2);
  for (int i = 0; i < N; i++) {
    if (~color[i])
      continue;
    vertex.push_back(i);
  }
  int ret = 114514, beet;
  for (auto &i : vertex) {
    int cost = v[0][i].first + v[1][i].first + v[2][i].first;
    int latte = min({v[0][i].second, v[1][i].second, v[2][i].second});
    if (cost < ret)
      ret = cost, beet = latte;
    else if (cost == ret)
      beet = min(beet, latte);
  }
  printf("%d %d\\n", ret, beet + 1);
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

const int SHIFT = 10000;
const int INF = 1 << 29;

int N, A, B, C, M;
int color[10000], small[10005];
vector<int> g[10005];
pair<int, int> v[10005];

pair<int, int> solve(int start) {
  fill_n(v, 10005, make_pair(INF, INF));

  queue<int> que;
  v[start] = {0, small[start]};
  que.emplace(start);
  while (!que.empty()) {
    int p = que.front();
    que.pop();
    for (auto &to : g[p]) {
      if (v[p].first + 1 == v[to].first) {
        v[to].second = min(v[to].second, v[p].second);
      } else if (v[p].first + 1 < v[to].first) {
        v[to] = {v[p].first + 1, min(small[to], v[p].second)};
        que.emplace(to);
      }
    }
  }
  return {v[SHIFT].first + v[SHIFT + 1].first + v[SHIFT + 2].first,
          min({v[SHIFT].second, v[SHIFT + 1].second, v[SHIFT + 2].second})};
}

int main() {
  memset(color, -1, sizeof(color));
  fill_n(small, 10005, INF);

  scanf("%d %d %d %d", &N, &A, &B, &C);
  for (int i = 0; i < N; i++)
    small[i] = i;
  for (int i = 0; i < A + B + C; i++) {
    int K;
    scanf("%d", &K);
    --K;
    if (i < A)
      color[K] = 0;
    else if (i < A + B)
      color[K] = 1;
    else
      color[K] = 2;
    small[color[K] + SHIFT] = min(small[color[K] + SHIFT], K);
  }

  scanf("%d", &M);
  while (M--) {
    int a, b;
    scanf("%d %d", &a, &b);
    --a, --b;
    if (~color[a])
      a = color[a] + SHIFT;
    if (~color[b])
      b = color[b] + SHIFT;
    g[a].emplace_back(b);
    g[b].emplace_back(a);
  }
  for (int i = 0; i < 10005; i++) {
    sort(begin(g[i]), end(g[i]));
    g[i].erase(unique(begin(g[i]), end(g[i])), end(g[i]));
  }

  vector<int> vertex;
  vertex.push_back(SHIFT);
  vertex.push_back(SHIFT + 1);
  vertex.push_back(SHIFT + 2);
  for (int i = 0; i < N; i++) {
    if (~color[i])
      continue;
    vertex.push_back(i);
  }

  int ret = 114514, beet;
  int ptr = 0;
  reverse(begin(vertex), end(vertex));
  for (int i : vertex) {
    auto get = solve(i);
    ++ptr;
    if (get.first < ret)
      ret = get.first, beet = get.second;
    else if (get.first == ret)
      beet = min(beet, get.second);
    if (ptr > 8000)
      break;
  }
  printf("%d %d\\n", ret, beet + 1);
}