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

#include <algorithm>
#include <bitset>
#include <cmath>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <vector>
#define REP(i, n) for (int i = 0; i < n; i++)
#define INF 1000000007
using namespace std;
using pii = pair<int, int>;
using LL = int64_t;

int main() {
  int n, k;
  cin >> n >> k;
  vector<pair<LL, int>> a(n);
  vector<LL> neg, pos;
  int pos_cnt = 0, neg_cnt = 0;
  LL ans = 1;

  REP(i, n) {
    LL b;
    cin >> b;
    if (b >= 0) {
      a.at(i) = make_pair(b, 1);
      pos.push_back(b);
      pos_cnt++;
    } else {
      a.at(i) = make_pair(-b, -1);
      neg.push_back(-b);
      neg_cnt++;
    }
  }

  if (k == n) {

    REP(i, n) {
      ans *= a.at(i).first * a.at(i).second;
      ans %= INF;
    }
    if (ans < 0) {
      ans += INF;
      ans %= INF;
    }

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

  if (neg_cnt == n && k % 2 == 1) {
    sort(a.begin(), a.end());

    REP(i, k) {
      ans *= a.at(i).first;
      ans %= INF;
    }
    ans *= -1;
    ans %= INF;
    ans += INF;
    ans %= INF;
    cout << ans << endl;
    return 0;
  }

  sort(a.begin(), a.end(), greater<pair<LL, int>>());
  sort(pos.begin(), pos.end(), greater<LL>());
  sort(neg.begin(), neg.end(), greater<LL>());

  int pos_number = 0, neg_number = 0;

  REP(i, k) {

    if (a.at(i).second > 0) {
      pos_number++;
    } else {
      neg_number++;
    }
  }

  if (neg_number % 2 == 0) {
    REP(i, k) {
      ans *= a.at(i).first;
      ans %= INF;
    }
    cout << ans << endl;
    return 0;
  }

  // 本番
  int i = pos_number, j = neg_number; // i+j=k<n,j:odd>0

  if (i == 0) {
    REP(k, i + 1) {
      ans *= pos[k];
      ans %= INF;
    }
    REP(k, j - 1) {
      ans *= neg[k];
      ans %= INF;
    }
    cout << ans << endl;
    return 0;
  }

  // i>0,j>0,i+j=k<n, j:odd
  if (i < pos_cnt && j < neg_cnt) {
    LL x = pos[i] * pos[i - 1], y = neg[j] * neg[j - 1];

    if (y > x) {

      REP(k, i - 1) {
        ans *= pos[k];
        ans %= INF;
      }
      REP(k, j + 1) {
        ans *= neg[k];
        ans %= INF;
      }
      cout << ans << endl;
      return 0;
    } else {

      REP(k, i + 1) {
        ans *= pos[k];
        ans %= INF;
      }
      REP(k, j - 1) {
        ans *= neg[k];
        ans %= INF;
      }
      cout << ans << endl;
      return 0;
    }
  }

  if (j == neg_cnt) {
    REP(k, i + 1) {
      ans *= pos[k];
      ans %= INF;
    }
    REP(k, j - 1) {
      ans *= neg[k];
      ans %= INF;
    }
    cout << ans << endl;
    return 0;
  }

  REP(k, i - 1) {
    ans *= pos[k];
    ans %= INF;
  }
  REP(k, j + 1) {
    ans *= neg[k];
    ans %= INF;
  }
  cout << ans << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cmath>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <vector>
#define REP(i, n) for (int i = 0; i < n; i++)
#define INF 1000000007
using namespace std;
using pii = pair<int, int>;
using LL = int64_t;

int main() {
  int n, k;
  cin >> n >> k;
  vector<pair<int, int>> a(n);
  vector<int> neg, pos;
  int pos_cnt = 0, neg_cnt = 0;
  LL ans = 1;

  REP(i, n) {
    int b;
    cin >> b;
    if (b >= 0) {
      a.at(i) = make_pair(b, 1);
      pos.push_back(b);
      pos_cnt++;
    } else {
      a.at(i) = make_pair(-b, -1);
      neg.push_back(-b);
      neg_cnt++;
    }
  }

  if (k == n) {

    REP(i, n) {
      ans *= a.at(i).first * a.at(i).second;
      ans %= INF;
    }
    if (ans < 0) {
      ans += INF;
      ans %= INF;
    }

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

  if (neg_cnt == n && k % 2 == 1) {
    sort(a.begin(), a.end());

    REP(i, k) {
      ans *= a.at(i).first;
      ans %= INF;
    }
    ans *= -1;
    ans %= INF;
    ans += INF;
    ans %= INF;
    cout << ans << endl;
    return 0;
  }

  sort(a.begin(), a.end(), greater<pair<int, int>>());
  sort(pos.begin(), pos.end(), greater<int>());
  sort(neg.begin(), neg.end(), greater<int>());

  int pos_number = 0, neg_number = 0;

  REP(i, k) {

    if (a.at(i).second > 0) {
      pos_number++;
    } else {
      neg_number++;
    }
  }

  if (neg_number % 2 == 0) {
    REP(i, k) {
      ans *= a.at(i).first;
      ans %= INF;
    }
    cout << ans << endl;
    return 0;
  }

  // 本番
  int i = pos_number, j = neg_number; // i+j=k<n,j:odd>0

  if (i == 0) {
    REP(k, i + 1) {
      ans *= pos[k];
      ans %= INF;
    }
    REP(k, j - 1) {
      ans *= neg[k];
      ans %= INF;
    }
    cout << ans << endl;
    return 0;
  }

  // i>0,j>0,i+j=k<n, j:odd
  if (i < pos_cnt && j < neg_cnt) {
    LL x = pos[i] * pos[i - 1], y = neg[j] * neg[j - 1];

    if (y >= x) {
      while (true) {
      }
      REP(k, i - 1) {
        ans *= pos[k];
        ans %= INF;
      }
      REP(k, j + 1) {
        ans *= neg[k];
        ans %= INF;
      }
      cout << ans << endl;
      return 0;
    } else {
      REP(k, i + 1) {
        ans *= pos[k];
        ans %= INF;
      }
      REP(k, j - 1) {
        ans *= neg[k];
        ans %= INF;
      }
      cout << ans << endl;
      return 0;
    }
  }

  /*if(j==neg_cnt){
      REP(k,i+1){
          ans*=pos[k];
          ans%=INF;
      }
      REP(k,j-1){
          ans*=neg[k];
          ans%=INF;
      }
      cout<<ans<<endl;
      return 0;
  }*/

  /*REP(k,i-1){
          ans*=pos[k];
          ans%=INF;
      }
      REP(k,j+1){
          ans*=neg[k];
          ans%=INF;
      }
      cout<<ans<<endl;*/

  while (true) {
  }

  return 0;
}