Strorkisのブログ

信憑性については悪しからず

ABC155D

問題

D - Pairs

 N個の整数 A_iにおける全てのペアの積を求め、それらを小さい順に並び替えたとき、 K番目にくる数を求めよ。

解法

「積が ans以下であるペアの個数が K個」となる ansの最小値を、二分探索で求める。

最左値と右辺値は入力された値から求めても良いが、正負によって場合分けしなければいけないため、 A_iの範囲にしておくのが楽。

計算量は、 O(log_{2}(2*10^{9}))


積が ans以下であるペアの個数は、積が負・正・0の場合に分けて求める。

まず、 A_{i}を負・正・0に分ける。

積が負の場合、正負のどちらかのリストを絶対値が大きい方から、もう一方のリストを絶対値が小さいほうから値を取り出すことで、尺取り法をすることができる。

積が正の場合、正負のリスト共に、絶対値が大きい方からと小さい方から値を取り出すことで、尺取り法をすることができる。これは、ペアの片方を固定したとき、もう片方を固定した値よりも小さい値から選べば良いという考えに基づいている。

積が0の場合、積が負になる個数と積が0になる個数を足せば良い。

計算量は積が0の場合を除き、 O(N)

コード

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
typedef long long ll;
#define MAXA 1000000000LL

int main() {
  ll n, k; cin >> n >> k;
  vector<ll> vm, vp;
  ll zero = 0;
  for (int i = 0; i < n; i++) {
    ll a; cin >> a;
    if (a < 0)
      vm.push_back(a);
    else if (a > 0)
      vp.push_back(a);
    else
      zero++;
  }
  sort(vm.begin(), vm.end());
  sort(vp.begin(), vp.end(), greater<ll>());

  auto f1 = [](vector<ll>& v1, vector<ll>& v2, ll x) {
    ll res = 0;
    int i = 0, j = v2.size()-1;
    while (i < v1.size() && j >= 0) {
      if (v1[i] * v2[j] <= x) {
        res += j + 1;
        i++;
      }
      else j--;
    }
    return res;
  };

  auto f2 = [](vector<ll>& v, ll x) {
    ll res = 0;
    int i = 0, j = v.size()-1;
    while (i < j) {
      if (v[i] * v[j] > x) {
        res += j - i;
        i++;
      }
      else j--;
    }
    return res;
  };

  ll way = n * (n-1) / 2;
  ll r = MAXA * MAXA, l = -r;
  while (r - l > 1) {
    ll c = (l + r) / 2;

    ll res;
    if (c < 0)
      res = f1(vm, vp, c);
    else if (c > 0)
      res = way - (f2(vm, c) + f2(vp, c));
    else {
      res = vm.size() * vp.size();
      res += (n-zero) * zero + zero * (zero-1) / 2;
    }

    if (res < k)
      l = c;
    else
      r = c;
  }
  cout << r << endl;
}

感想

二分探索と尺取り法だけでここまで考察が難しくなるのには驚かされた。

ただ、それだけに基本が詰まった面白い問題だったように思う。

(尺取り法の解説を書くのは難しい…)