Problem p02616 (C++)

Buggy submission (status: Runtime Error):


#include <algorithm>
#include <assert.h>
#include <bitset>
#include <iostream>
#include <map>
#include <math.h>
#include <memory.h>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <stdio.h>
#include <string>
#include <time.h>
#include <unordered_map>
#include <vector>

using namespace std;

const int mod = 1000000007;

int a[111111];

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

int main() {
  ios::sync_with_stdio(0);
  int t, sn = 0;
  t = 1; // assert(1 <= t && t <= 1e5);
  while (t--) {
    int n, k;
    cin >> n >> k;
    sn += n;
    bool zero = 0;
    for (int i = 0; i < n; ++i) {
      cin >> a[i];

      zero |= (a[i] == 0);
    }
    sort(a, a + n, cmp);

    //        for(int i = 0; i < n; ++i) {
    //            cout << a[i] << " ";
    //            if (i == n-k-1)
    //                cout << "\n";
    //        }
    //        cout << "\n";

    int sn = 1;
    vector<int> nums;
    int gneg = -1e9 - 1, gpos = 1e9 + 1;
    for (int i = 0, id = n - 1; i < k; ++i, id--) {
      if (a[id] < 0)
        sn *= -1;
      nums.push_back(a[id]);
      if (a[id] < 0)
        gneg = max(gneg, a[id]);
      else
        gpos = min(gpos, a[id]);
    }

    if (sn > 0) {
      int ans = 1;
      for (auto x : nums)
        ans = (ans * 1ll * x) % mod;
      cout << (ans + mod) % mod << "\n";
    } else {
      int mn = 1, mx = -1;
      for (int i = 0; i < n - k; ++i) {
        mn = min(mn, a[i]);
        mx = max(mx, a[i]);
      }

      // cout << gneg << " " << mx << "\n";
      // cout << gpos << " " << mn << "\n";

      bool g1 = gneg != -1e-9 - 1 && mx != -1;
      bool g2 = gpos != 1e9 + 1 && mn != 1;
      if (g1 && g2) {
        int val, repl;
        if (mx * 1ll * gpos > mn * 1ll * gneg)
          val = gneg, repl = mx;
        else
          val = gpos, repl = mn;

        for (int i = 0; i < k; ++i)
          if (nums[i] == val) {
            nums[i] = repl;
            break;
          }
        int ans = 1;
        for (auto x : nums)
          ans = (ans * 1ll * x) % mod;
        cout << (ans + mod) % mod << "\n";
      } else if (g1) {
        for (int i = 0; i < k; ++i)
          if (nums[i] == gneg) {
            nums[i] = mx;
            break;
          }
        int ans = 1;
        for (auto x : nums)
          ans = (ans * 1ll * x) % mod;
        cout << (ans + mod) % mod << "\n";
      } else if (g2) {
        for (int i = 0; i < k; ++i)
          if (nums[i] == gpos) {
            nums[i] = mn;
            break;
          }
        int ans = 1;
        for (auto x : nums)
          ans = (ans * 1ll * x) % mod;
        cout << (ans + mod) % mod << "\n";
      } else {
        if (zero)
          cout << "0\n";
        else {
          int ans = 1;
          for (int i = 0; i < k; ++i)
            ans = (ans * 1ll * a[i]) % mod;
          cout << (ans + mod) % mod << "\n";
        }
      }
    }
  }
  return 0;
}

Fixed submission:


#include <algorithm>
#include <assert.h>
#include <bitset>
#include <iostream>
#include <map>
#include <math.h>
#include <memory.h>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <stdio.h>
#include <string>
#include <time.h>
#include <unordered_map>
#include <vector>

using namespace std;

const int mod = 1000000007;

int a[211111];

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

int main() {
  ios::sync_with_stdio(0);
  int t, sn = 0;
  t = 1; // assert(1 <= t && t <= 1e5);
  while (t--) {
    int n, k;
    cin >> n >> k;
    sn += n;
    bool zero = 0;
    for (int i = 0; i < n; ++i) {
      cin >> a[i];

      zero |= (a[i] == 0);
    }
    sort(a, a + n, cmp);

    //        for(int i = 0; i < n; ++i) {
    //            cout << a[i] << " ";
    //            if (i == n-k-1)
    //                cout << "\n";
    //        }
    //        cout << "\n";

    int sn = 1;
    vector<int> nums;
    int gneg = -1e9 - 1, gpos = 1e9 + 1;
    for (int i = 0, id = n - 1; i < k; ++i, id--) {
      if (a[id] < 0)
        sn *= -1;
      nums.push_back(a[id]);
      if (a[id] < 0)
        gneg = max(gneg, a[id]);
      else
        gpos = min(gpos, a[id]);
    }

    if (sn > 0) {
      int ans = 1;
      for (auto x : nums)
        ans = (ans * 1ll * x) % mod;
      cout << (ans + mod) % mod << "\n";
    } else {
      int mn = 1, mx = -1;
      for (int i = 0; i < n - k; ++i) {
        mn = min(mn, a[i]);
        mx = max(mx, a[i]);
      }

      // cout << gneg << " " << mx << "\n";
      // cout << gpos << " " << mn << "\n";

      bool g1 = gneg != -1e-9 - 1 && mx != -1;
      bool g2 = gpos != 1e9 + 1 && mn != 1;
      if (g1 && g2) {
        int val, repl;
        if (mx * 1ll * gpos > mn * 1ll * gneg)
          val = gneg, repl = mx;
        else
          val = gpos, repl = mn;

        for (int i = 0; i < k; ++i)
          if (nums[i] == val) {
            nums[i] = repl;
            break;
          }
        int ans = 1;
        for (auto x : nums)
          ans = (ans * 1ll * x) % mod;
        cout << (ans + mod) % mod << "\n";
      } else if (g1) {
        for (int i = 0; i < k; ++i)
          if (nums[i] == gneg) {
            nums[i] = mx;
            break;
          }
        int ans = 1;
        for (auto x : nums)
          ans = (ans * 1ll * x) % mod;
        cout << (ans + mod) % mod << "\n";
      } else if (g2) {
        for (int i = 0; i < k; ++i)
          if (nums[i] == gpos) {
            nums[i] = mn;
            break;
          }
        int ans = 1;
        for (auto x : nums)
          ans = (ans * 1ll * x) % mod;
        cout << (ans + mod) % mod << "\n";
      } else {
        if (zero)
          cout << "0\n";
        else {
          int ans = 1;
          for (int i = 0; i < k; ++i)
            ans = (ans * 1ll * a[i]) % mod;
          cout << (ans + mod) % mod << "\n";
        }
      }
    }
  }
  return 0;
}