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

#include <bits/stdc++.h>
#define fo(i, a, b) for (int i = a; i <= b; i++)
#define fod(i, a, b) for (int i = a; i >= b; i--)
#define me0(a) memset(a, 0, sizeof(a))
#define me1(a) memset(a, -1, sizeof(a))
#define op freopen("in.txt", "r", stdin)
#define pii pair<int, int>
#define mp(x, y) make_pair(x, y)
using namespace std;
const int INF = 0x3f3f3f3f;
typedef long long LL;
void read(int &val) {
  int x = 0;
  int bz = 1;
  char c;
  for (c = getchar(); (c < '0' || c > '9') && c != '-'; c = getchar())
    ;
  if (c == '-') {
    bz = -1;
    c = getchar();
  }
  for (; c >= '0' && c <= '9'; c = getchar())
    x = x * 10 + c - 48;
  val = x * bz;
}
const int mod = 1e9 + 7;
const int maxn = 2e6 + 1000;

int n, a[maxn], k;

LL ksm(LL a, LL b) {
  a = (a % mod + mod) % mod;
  LL res = 1;
  while (b) {
    if (b & 1) {
      res = res * a % mod;
    }
    a = a * a % mod;
    b >>= 1;
  }
  return res;
}
bool cmp1(int a, int b) { return abs(a) > abs(b); }
bool cmp2(int a, int b) { return abs(a) < abs(b); }

int main() {
  //    op;
  read(n);
  read(k);
  int num1 = 0, num0 = 0, num2 = 0;
  fo(i, 1, n) {
    read(a[i]);
    if (a[i] > 0)
      num2++;
    else if (a[i] == 0)
      num0++;
    else
      num1++;
  }
  LL ans = 1;
  if (num0 > n - k) {
    puts("0");
    return 0;
  }
  if (k == n) {
    fo(i, 1, n) ans = ans * a[i] % mod;
  } else if (num2 == 0 && (k & 1)) {
    sort(a + 1, a + 1 + n, cmp2);
    fo(i, 1, k) ans = ans * a[i] % mod;
  } else {
    int cnt1 = 0;
    sort(a + 1, a + 1 + n, cmp1);
    fo(i, 1, k) {
      if (a[i] < 0)
        cnt1++;
      ans = ans * a[i] % mod;
    }
    if (cnt1 & 1) {
      if (cnt1 == k) {
        ans = ans * ksm(a[k], mod - 2) % mod;
        fo(i, k + 1, n) {
          if (a[i] >= 0) {
            ans = ans * a[i] % mod;
            break;
          }
        }
      } else {
        LL t1 = 1, t2 = 1, flag1 = 1, flag2 = 1;
        fo(i, 1, k) {
          if (a[i] < 0)
            t1 = a[i];
          else
            t2 = a[i];
        }

        ans = ans * ksm(t1 * t2 % mod, mod - 2) % mod;

        fo(i, k + 1, n) {
          if (flag1 && a[i] < 0) {
            t1 *= a[i];
            flag1 = 0;
          } else if (flag2 && a[i] >= 0) {
            t2 *= a[i];
            flag2 = 0;
          }
        }
        ans = ans * (max(t1, t2) % mod) % mod;
      }
    }
  }
  ans = (ans % mod + mod) % mod;
  printf("%lld\\n", ans);

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define fo(i, a, b) for (int i = a; i <= b; i++)
#define fod(i, a, b) for (int i = a; i >= b; i--)
#define me0(a) memset(a, 0, sizeof(a))
#define me1(a) memset(a, -1, sizeof(a))
#define op freopen("in.txt", "r", stdin)
#define pii pair<int, int>
#define mp(x, y) make_pair(x, y)
using namespace std;
const int INF = 0x3f3f3f3f;
typedef unsigned long long ULL;
typedef long long LL;
void read(LL &val) {
  int x = 0;
  int bz = 1;
  char c;
  for (c = getchar(); (c < '0' || c > '9') && c != '-'; c = getchar())
    ;
  if (c == '-') {
    bz = -1;
    c = getchar();
  }
  for (; c >= '0' && c <= '9'; c = getchar())
    x = x * 10 + c - 48;
  val = x * bz;
}
const int mod = 1e9 + 7;
const int maxn = 1e6 + 11;
LL n, m, k;

LL a[maxn], b[maxn], c[maxn], d[maxn];
int tot1 = 0, tot2 = 0;
void cal() {
  sort(b + 1, b + 1 + n);
  LL res = 1;
  fo(i, 1, k) {
    res *= b[i];
    res %= mod;
  }
  if (res)
    res = -res;
  res = (res % mod + mod) % mod;
  printf("%lld\\n", res);
}

bool cmp(LL a, LL b) { return abs(a) > abs(b); }

LL ksm(LL a, LL b) {
  LL res = 1;
  while (b) {
    if (b & 1) {
      res = res * a % mod;
    }
    a = a * a % mod;
    b >>= 1;
  }
  res = (res + mod) % mod;
  return res;
}

int main() {
  op;
  read(n);
  read(k);
  int num0 = 0;
  fo(i, 1, n) {
    read(a[i]);
    b[i] = abs(a[i]);
    if (a[i] < 0)
      c[++tot1] = a[i];
    else if (a[i] > 0)
      d[++tot2] = a[i];
    else
      num0++;
  }

  if (n == k) {
    LL res = 1;
    fo(i, 1, n) { res = res * a[i] % mod; }
    res = (res % mod + mod) % mod;
    printf("%lld\\n", res);
    return 0;
  }
  if (num0 > n - k) {
    puts("0");
    return 0;
  }
  if ((tot2 == 0) && (k & 1)) {
    cal();
    return 0;
  }

  LL res = 1;
  sort(a + 1, a + 1 + n, cmp);

  LL cnt1 = 0, mx = -INF, mx2 = INF, cnt2 = 0;

  fo(i, 1, k) {
    if (a[i] < 0) {
      cnt1++;
      mx = a[i];
    } else {
      mx2 = a[i];
    }
    res = res * a[i] % mod;
  }
  if ((cnt1 & 1) && (cnt1 == k)) {
    res = res * ksm(mx, mod - 2) % mod;
    fo(i, k + 1, n) {
      if (a[i] >= 0) {
        res *= a[i];
        res %= mod;
        cnt2++;
        if (cnt2 >= 2)
          break;
      }
    }
  } else if (cnt1 & 1) {
    LL t1 = mx, t2 = mx2; // 负正
    int flag1 = 1, flag2 = 1;

    fo(i, k + 1, n) {
      if (mx2 != INF && a[i] >= 0 && flag2) {
        flag2 = 0;
        t2 *= a[i];
      } else if (a[i] < 0 && flag1) {
        flag1 = 0;
        t1 *= a[i];
      }
    }
    res = (res * ksm(mx * mx2 % mod, mod - 2) % mod + mod) % mod;

    if (!flag1 && !flag2) {
      if (t1 > t2) {
        res = res * t1 % mod;
      } else {
        res = res * t2 % mod;
      }
    } else if (!flag1) {
      res = res * t1 % mod;
    } else if (!flag2) {
      res = res * t2 % mod;
    } else {
      res = res * mx % mod * mx2 % mod;
    }
  }
  res = (res % mod + mod) % mod;
  printf("%lld\\n", res);
  return 0;
}