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

#include <algorithm>
#include <assert.h>
#include <climits>
#include <cmath>
#include <iostream>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using namespace std;
constexpr int64_t MOD = 1e9 + 7;

#include <unistd.h>

int main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);

  int32_t N, K;
  cin >> N >> K;
  vector<int64_t> p;
  vector<int64_t> m;
  vector<int64_t> z;
  for (int i = 0; i < N; i++) {
    int64_t a;
    cin >> a;
    if (a > 0) {
      p.push_back(a);
    } else if (a == 0) {
      z.push_back(a);
    } else {
      m.push_back(-a);
    }
  }
  sort(p.begin(), p.end(), greater<int64_t>());
  sort(m.begin(), m.end(), greater<int64_t>());
  int32_t np = p.size();
  int32_t nm = m.size();
  int32_t nz = z.size();
  int64_t res = 1;
  if (np == 0) {
    if (nm >= K && K % 2 == 0) {
      // nm big k
      for (int i = 0; i < K; i++) {
        res *= m[i];
        res %= MOD;
      }
    } else if (nz == 0) { // minus
      res = 1;
      sort(m.begin(), m.end());
      for (int i = 0; i < K; i++) {
        res *= -m[i];
        if (res < 0)
          res += MOD;
        res %= MOD;
      }
    } else {
      res = 0;
    }
  } else {
    if (np + nm < K) {
      res = 0;
    } else if (np + nm == K && nm % 2 == 1) {
      if (nz != 0) {
        res = 0;
      } else {
        // all np and nm
        for (int i = 0; i < np; i++) {
          res *= p[i];
          res %= MOD;
        }
        for (int i = 0; i < nm; i++) {
          res *= -m[i];
          res %= MOD;
          while (res < 0)
            res += MOD;
        }
      }
    } else { // np + nm > K, np > 0
      // take even number of m's
      if (K % 2 == 1) {
        int pidx = 1;
        int midx = 0;
        res = p[0];
        while (pidx + midx != K && pidx + 1 < np && midx + 1 < nm) {
          if (p[pidx] * p[pidx + 1] >= m[midx] * m[midx + 1]) {
            res *= p[pidx];
            res %= MOD;
            res *= p[pidx + 1];
            res %= MOD;
            pidx += 2;
          } else {
            res *= m[midx];
            res %= MOD;
            res *= m[midx + 1];
            res %= MOD;
            midx += 2;
          }
        }
        while (pidx + midx != K && pidx + 1 < np) {
          res *= p[pidx];
          res %= MOD;
          res *= p[pidx + 1];
          res %= MOD;
          pidx += 2;
        }
        while (pidx + midx != K && midx + 1 < nm) {
          res *= m[midx];
          res %= MOD;
          res *= m[midx + 1];
          res %= MOD;
          midx += 2;
        }
      } else {
        int pidx = 0;
        int midx = 0;
        res = 1;
        while (pidx + midx != K && pidx + 1 < np && midx + 1 < nm) {
          if (p[pidx] * p[pidx + 1] >= m[midx] * m[midx + 1]) {
            res *= p[pidx];
            res %= MOD;
            res *= p[pidx + 1];
            res %= MOD;
            pidx += 2;
          } else {
            res *= m[midx];
            res %= MOD;
            res *= m[midx + 1];
            res %= MOD;
            midx += 2;
          }
        }
        while (pidx + midx != K && pidx + 1 < np) {
          res *= p[pidx];
          res %= MOD;
          res *= p[pidx + 1];
          res %= MOD;
          pidx += 2;
        }
        while (pidx + midx != K && midx + 1 < nm) {
          res *= m[midx];
          res %= MOD;
          res *= m[midx + 1];
          res %= MOD;
          midx += 2;
        }
      }
    }
  }
  cout << res << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <assert.h>
#include <climits>
#include <cmath>
#include <iostream>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using namespace std;
constexpr int64_t MOD = 1e9 + 7;

#include <unistd.h>

int main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);

  int32_t N, K;
  cin >> N >> K;
  vector<int64_t> p;
  vector<int64_t> m;
  vector<int64_t> z;
  for (int i = 0; i < N; i++) {
    int64_t a;
    cin >> a;
    if (a > 0) {
      p.push_back(a);
    } else if (a == 0) {
      z.push_back(a);
    } else {
      m.push_back(-a);
    }
  }
  sort(p.begin(), p.end(), greater<int64_t>());
  sort(m.begin(), m.end(), greater<int64_t>());
  int32_t np = p.size();
  int32_t nm = m.size();
  int32_t nz = z.size();
  int64_t res = 1;
  if (np == 0) {
    if (nm >= K && K % 2 == 0) {
      // nm big k
      for (int i = 0; i < K; i++) {
        res *= m[i];
        res %= MOD;
      }
    } else if (nz == 0) { // minus
      res = 1;
      sort(m.begin(), m.end());
      for (int i = 0; i < K; i++) {
        res *= -m[i];
        if (res < 0)
          res += MOD;
        res %= MOD;
      }
    } else {
      res = 0;
    }
  } else {
    if (np + nm < K) {
      res = 0;
    } else if (np + nm == K && nm % 2 == 1) {
      if (nz != 0) {
        res = 0;
      } else {
        // all np and nm
        for (int i = 0; i < np; i++) {
          res *= p[i];
          res %= MOD;
        }
        for (int i = 0; i < nm; i++) {
          res *= -m[i];
          while (res < 0)
            res += MOD;
          res %= MOD;
        }
      }
    } else { // np + nm > K, np > 0
      // take even number of m's
      if (K % 2 == 1) {
        int pidx = 1;
        int midx = 0;
        res = p[0];
        while (pidx + midx != K && pidx + 1 < np && midx + 1 < nm) {
          if (p[pidx] * p[pidx + 1] >= m[midx] * m[midx + 1]) {
            res *= p[pidx];
            res %= MOD;
            res *= p[pidx + 1];
            res %= MOD;
            pidx += 2;
          } else {
            res *= m[midx];
            res %= MOD;
            res *= m[midx + 1];
            res %= MOD;
            midx += 2;
          }
        }
        while (pidx + midx != K && pidx + 1 < np) {
          res *= p[pidx];
          res %= MOD;
          res *= p[pidx + 1];
          res %= MOD;
          pidx += 2;
        }
        while (pidx + midx != K && midx + 1 < nm) {
          res *= m[midx];
          res %= MOD;
          res *= m[midx + 1];
          res %= MOD;
          midx += 2;
        }
      } else {
        int pidx = 0;
        int midx = 0;
        res = 1;
        while (pidx + midx != K && pidx + 1 < np && midx + 1 < nm) {
          if (p[pidx] * p[pidx + 1] >= m[midx] * m[midx + 1]) {
            res *= p[pidx];
            res %= MOD;
            res *= p[pidx + 1];
            res %= MOD;
            pidx += 2;
          } else {
            res *= m[midx];
            res %= MOD;
            res *= m[midx + 1];
            res %= MOD;
            midx += 2;
          }
        }
        while (pidx + midx != K && pidx + 1 < np) {
          res *= p[pidx];
          res %= MOD;
          res *= p[pidx + 1];
          res %= MOD;
          pidx += 2;
        }
        while (pidx + midx != K && midx + 1 < nm) {
          res *= m[midx];
          res %= MOD;
          res *= m[midx + 1];
          res %= MOD;
          midx += 2;
        }
      }
    }
  }
  cout << res << endl;
  return 0;
}