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

/*{{{*/
#include <bits/stdc++.h>

typedef long long ll;

#define REP(i, n) for (ll i = 0; i < n; i++)
#define FOR(i, a, b) for (ll i = a; i < b; i++)
#define FORR(i, a, b) for (ll i = b - 1; i >= a; i--)

#define vl vector<ll>
#define pll pair<ll, ll>

#define DEBUG(x) cout << #x << ": " << x << endl

using namespace std;
/*}}}*/

ll MOD = 1e9 + 7;
int main() {
  ll n, k;
  cin >> n >> k;
  vl pos, neg;
  REP(i, n) {
    ll a;
    cin >> a;
    if (a > 0)
      pos.push_back(a);
    if (a < 0)
      neg.push_back(-a);
  }
  bool exists0 = pos.size() + neg.size() < n;
  // sort by abs decending
  sort(pos.begin(), pos.end());
  reverse(pos.begin(), pos.end());
  sort(neg.begin(), neg.end());
  reverse(neg.begin(), neg.end());

  // when answer is 0
  if (pos.size() + neg.size() < k) {
    // always contains 0
    cout << 0 << endl;
    return 0;
  }
  if (exists0 && pos.size() == 0 && k % 2 == 1) {
    // contains 0, otherwise ans is negative
    cout << 0 << endl;
    return 0;
  }

  // when answer is negative
  if (!exists0 && pos.size() == 0 && k % 2 == 1) {
    reverse(neg.begin(), neg.end());
    ll ans = 1;
    REP(i, k) {
      ans *= neg[i];
      ans %= MOD;
    }
    ans = MOD - ans;
    cout << ans << endl;
    return 0;
  }
  if (!exists0 && k == n && neg.size() % 2 == 1) {
    // select all non-zero variables
    ll ans = 1;
    REP(i, neg.size()) {
      ans *= neg[i];
      ans %= MOD;
    }
    REP(i, pos.size()) {
      ans *= pos[i];
      ans %= MOD;
    }
    ans = MOD - ans;
    cout << ans << endl;
    return 0;
  }

  // when answer is positive
  ll cp = 0;
  ll cn = 0;
  if (pos.size() == 0)
    cn = k;
  else if (neg.size() == 0)
    cp = k;
  else
    REP(i, k) {
      if (cn < neg.size() && pos[cp] < neg[cn]) {
        cn++;
      } else {
        cp++;
      }
    }
  // num of employed negative is odd
  if (cn % 2 != 0) {
    if (cp == 0 && pos.size() > 0) {
      cp++;
      cn--;
    } else if (cp < pos.size() && cn < neg.size()) {
      ll candp = pos[cp - 1] * pos[cp];
      ll candn = neg[cn - 1] * neg[cn];
      if (candp > candn) {
        cp = cp + 1;
        cn = cn - 1;
      } else {
        cp = cp - 1;
        cn = cn + 1;
      }
    } else if (cn == neg.size()) {
      cp = cp + 1;
      cn = cn - 1;
    } else if (cp == pos.size()) {
      cp = cp - 1;
      cn = cn + 1;
    }
  }

  ll ans = 1;
  REP(i, cp) {
    ans *= pos[i];
    ans %= MOD;
  }
  REP(i, cn) {
    ans *= neg[i];
    ans %= MOD;
  }

  cout << ans << endl;
  return 0;
}


Error: TLE

Fixed submission:

/*{{{*/
#include <bits/stdc++.h>

typedef long long ll;

#define REP(i, n) for (ll i = 0; i < n; i++)
#define FOR(i, a, b) for (ll i = a; i < b; i++)
#define FORR(i, a, b) for (ll i = b - 1; i >= a; i--)

#define vl vector<ll>
#define pll pair<ll, ll>

#define DEBUG(x) cout << #x << ": " << x << endl

using namespace std;
/*}}}*/

ll MOD = 1e9 + 7;
int main() {
  ll n, k;
  cin >> n >> k;
  vl pos, neg;
  REP(i, n) {
    ll a;
    cin >> a;
    if (a > 0)
      pos.push_back(a);
    if (a < 0)
      neg.push_back(-a);
  }
  bool exists0 = pos.size() + neg.size() < n;
  // sort by abs decending
  sort(pos.begin(), pos.end());
  reverse(pos.begin(), pos.end());
  sort(neg.begin(), neg.end());
  reverse(neg.begin(), neg.end());

  // when answer is 0
  if (pos.size() + neg.size() < k) {
    // always contains 0
    cout << 0 << endl;
    return 0;
  }
  if (exists0 && pos.size() == 0 && k % 2 == 1) {
    // contains 0, otherwise ans is negative
    cout << 0 << endl;
    return 0;
  }

  // when answer is negative
  if (!exists0 && pos.size() == 0 && k % 2 == 1) {
    reverse(neg.begin(), neg.end());
    ll ans = 1;
    REP(i, k) {
      ans *= neg[i];
      ans %= MOD;
    }
    ans = MOD - ans;
    cout << ans << endl;
    return 0;
  }
  if (!exists0 && k == n && neg.size() % 2 == 1) {
    // select all non-zero variables
    ll ans = 1;
    REP(i, neg.size()) {
      ans *= neg[i];
      ans %= MOD;
    }
    REP(i, pos.size()) {
      ans *= pos[i];
      ans %= MOD;
    }
    ans = MOD - ans;
    cout << ans << endl;
    return 0;
  }
  ll ans = 1;
  REP(i, 1e12) {
    ans += i;
    ans %= MOD;
  }

  // when answer is positive
  ll cp = 0;
  ll cn = 0;
  if (pos.size() == 0)
    cn = k;
  else if (neg.size() == 0)
    cp = k;
  else
    REP(i, k) {
      if (cn < neg.size() && pos[cp] < neg[cn]) {
        cn++;
      } else {
        cp++;
      }
    }
  // num of employed negative is odd
  if (cn % 2 != 0) {
    // here cn >= 1 && cp >= 1
    if (cp < pos.size() && cn < neg.size()) {
      ll candp = pos[cp - 1] * pos[cp];
      ll candn = neg[cn - 1] * neg[cn];
      if (candp > candn) {
        cp = cp + 1;
        cn = cn - 1;
      } else {
        cp = cp - 1;
        cn = cn + 1;
      }
    } else if (cn == neg.size()) {
      cp = cp + 1;
      cn = cn - 1;
    } else if (cp == pos.size()) {
      cp = cp - 1;
      cn = cn + 1;
    }
  }

  // ll ans = 1;
  REP(i, cp) {
    ans *= pos[i];
    ans %= MOD;
  }
  REP(i, cn) {
    ans *= neg[i];
    ans %= MOD;
  }

  cout << ans << endl;
  return 0;
}