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

#include <iostream>
#include <set>
#include <stack>
#include <vector>

using namespace std;

int make_ballstack(vector<int> &balls) {
  int n = (int)balls.size();
  stack<pair<int, int>> stk_ball;
  int length;
  int ans = 0;

  for (int i = 0; i < n; ++i) {
    if (stk_ball.empty()) {
      balls[i];
      length = 1;
      ans = 1;
      stk_ball.push(make_pair(balls[i], length));
    } else if (stk_ball.top().first == balls[i]) {
      length = stk_ball.top().second;
      stk_ball.pop();
      ++length;
      ++ans;
      stk_ball.push(make_pair(balls[i], length));
    } else {
      length = stk_ball.top().second;
      if (length >= 4) {
        ans -= length;
        stk_ball.pop();
        if (stk_ball.empty()) {
          stk_ball.push(make_pair(balls[i], 1));
        } else if (stk_ball.top().first != balls[i]) {
          stk_ball.push(make_pair(balls[i], 1));
        } else {
          length = stk_ball.top().second;
          stk_ball.pop();
          stk_ball.push(make_pair(balls[i], length + 1));
        }
      } else {
        stk_ball.push(make_pair(balls[i], 1));
      }
      ++ans;
    }
  }
  if (stk_ball.top().second >= 4) {
    ans -= stk_ball.top().second;
    stk_ball.pop();
  }
  return ans;
}

int main() {
  int N;

  cin >> N;

  while (N != 0) {
    vector<int> balls;
    set<pair<int, int>> mark;
    int cont = 1;
    int temp;
    int ans = N;

    for (int i = 0; i < N; ++i) {
      cin >> temp;
      balls.push_back(temp);
      if (i >= 1 && temp == balls[i - 1])
        ++cont;
      else {
        if (cont == 3) {
          mark.insert(make_pair(i, balls[i - 1]));
          if (i >= 4) {
            mark.insert(make_pair(i - 4, balls[i - 1]));
          }
        }
        cont = 1;
      };
    }

    for (auto x : mark) {
      int tempc = balls[x.first];
      balls[x.first] = x.second;
      ans = min(ans, make_ballstack(balls));
      balls[x.first] = tempc;
    }

    cout << ans << endl;

    cin >> N;
  }

  return 0;
}

Error: TLE

Fixed submission:

#include <iostream>
#include <stack>
#include <vector>

using namespace std;

int make_ballstack(vector<int> &balls) {
  int n = (int)balls.size();
  stack<pair<int, int>> stk_ball;
  int length;
  int ans = 0;

  for (int i = 0; i < n; ++i) {
    if (stk_ball.empty()) {
      balls[i];
      length = 1;
      ans = 1;
      stk_ball.push(make_pair(balls[i], length));
    } else if (stk_ball.top().first == balls[i]) {
      length = stk_ball.top().second;
      stk_ball.pop();
      ++length;
      ++ans;
      stk_ball.push(make_pair(balls[i], length));
    } else {
      length = stk_ball.top().second;
      if (length >= 4) {
        ans -= length;
        stk_ball.pop();
        if (stk_ball.empty()) {
          stk_ball.push(make_pair(balls[i], 1));
        } else if (stk_ball.top().first != balls[i]) {
          stk_ball.push(make_pair(balls[i], 1));
        } else {
          length = stk_ball.top().second;
          stk_ball.pop();
          stk_ball.push(make_pair(balls[i], length + 1));
        }
      } else {
        stk_ball.push(make_pair(balls[i], 1));
      }
      ++ans;
    }
  }
  if (stk_ball.top().second >= 4) {
    ans -= stk_ball.top().second;
    stk_ball.pop();
  }
  return ans;
}

int main() {
  int N;

  cin >> N;

  while (N != 0) {
    vector<int> balls;
    int temp;
    int ans = N;

    for (int i = 0; i < N; ++i) {
      cin >> temp;
      balls.push_back(temp);
    }

    for (int i = 0; i < N; ++i) {
      balls[i] = balls[i] % 3 + 1;
      ans = min(ans, make_ballstack(balls));
      balls[i] = balls[i] % 3 + 1;
      ans = min(ans, make_ballstack(balls));
      balls[i] = balls[i] % 3 + 1;
    }

    cout << ans << endl;

    cin >> N;
  }

  return 0;
}