ABC193: E – Oversleeping を二分探索と Floor Sum で高速に解く

以下の問題を公式解説より良い計算量で解きます:
キャディプログラミングコンテスト2021 (AtCoder Beginner Contest 193): E - Oversleeping

想定解法は $O(YQ \log(X+Y+P+Q))$ や $O(Y + Q + \log(X+Y+P+Q))$ でしたが,この記事では $O(\log^2 (X+Y) + \log^2 (P+Q))$ で解きます.Floor Sum という道具を使います.

準備: Floor Sum とは?

Floor Sum とは以下の問題,もしくはその答えとなるアルゴリズムの通称です:

整数 $n, m, a, b$ が与えられる.ただし $m \geq 1$ である.このとき,\[ \sum_{k=0}^{n-1} \bigg\lfloor \frac{ak + b}{m} \bigg\rfloor \] を $O(\log m)$ 時間で計算せよ.

このアルゴリズムは,AtCoder Library にも入っています → ドキュメント

また,AtCoder Library Practice Contest でアルゴリズムが解説されています.これ自体もかなり面白いです → C - Floor Sum 解説

解法

2 通りの場合分け

高橋くんが街 $B$ で最短で降りられるのは,

  • 電車が街 $B$ に着いた瞬間
  • 高橋くんが起きた瞬間

のどちらかです.よって,それぞれの場合について最短時間を求め,その小さい方を答えればよいです.

どちらも本質的には同じなので,以下では「電車が街 $B$ に着いた瞬間」の場合のみ説明します.

問題の定式化

電車が街 $B$ に着く時刻 $t$ は,ある整数 $n \geq 0$ を用いて\begin{equation} \label{n2t} t = 2(X+Y)n + X \end{equation}と書けます.一方,高橋くんが起きている時刻 $t$ は,ある整数 $m \geq 0$ に対し,\[(P+Q)m + P \leq t < (P+Q)(m+1)\] を満たします.これらより,「不等式 \[ (P+Q)m + P \leq 2(X+Y)n + X < (P+Q)(m+1)\] を満たす整数 $m \geq 0$ が存在する」ような,整数 $n \geq 0$ の最小値を求めればよいです.

式変形と言い換え

「$m$ が存在するか?」に注目したいので,不等式を $m$ について解きます.すると,\[ \frac{2(X+Y)n + X - P - Q}{P+Q} < m \leq \frac{2(X+Y)n + X - P}{P+Q} \]となります.最左辺と最右辺をそれぞれ $L(n), R(n)$ とおきます.

今,「$L(n) < m \leq R(n)$ を満たす整数 $m$ が存在する」ような最小の $n$ を求めたいです.これを言い換えましょう.$R(n) - L(n) = \frac{Q}{P+Q} > 0$ なので,「$L(n) < m \leq R(n)$ を満たす整数 $m$」の個数 \[ \lfloor R(n) \rfloor - \lfloor L(n) \rfloor \ (\geq 0) \] です.つまり,$\lfloor R(n) \rfloor - \lfloor L(n) \rfloor > 0$ となる最小の $n$ を求めればよくなりました.

二分探索へ

$\lfloor R(n) \rfloor - \lfloor L(n) \rfloor$ は $n$ について単調増加ではありません.しかし,累積和をとった\[ f(n) := \sum_{k=0}^{n} \Big( \lfloor R(k) \rfloor - \lfloor L(k) \rfloor \Big) \]は $n$ について単調増加なので,$f(n) > 0$ となる最小の $n$ が二分探索で求められます.この $n$ が求めたかったものでした! 式 \eqref{n2t} より,$n$ から答えの時刻 $t$ が計算できます.

高橋くんが $(P+Q)$ 秒周期で行動することを考えると,二分探索の上限は $n = P+Q$ に設定すれば十分です.もし $f(P+Q-1) = 0$ ならば,高橋くんが起きた瞬間に街 $B$ にいることはありません.

$f(n) > 0$ かどうかは,Floor Sum を $2$ 回使うと $O(\log (P+Q))$ で判定できます.よって,この二分探索の計算量は $O(\log^2 (P+Q))$ です.

「高橋くんが起きた瞬間」も同様に考えると,解法全体の計算量は $1$ ケースあたり $O(\log^2 (X+Y) + \log^2 (P+Q))$ です.

オーバーフローしないか?

問題の答えや $\sum_{k=0}^{P+Q} \lfloor R(k) \rfloor$ の値は高々 $2(X+Y)(P+Q)$ 程度です.例えば $X, Y, P, Q \leq 10^9$ なら,これはギリギリ 64 bit 整数型に収まるので大丈夫です.

実装例 (C++ with ACL)

#include <bits/stdc++.h>
#include <atcoder/all>
using namespace std;
using namespace atcoder;
using ll = long long;
template<class T> inline bool chmin(T &a, const T b) { if (a > b) { a = b; return true; } return false; }
const ll INF = numeric_limits<ll>::max();;

ll subsolve(ll m, ll a, ll b, ll c1, ll c2) {
  ll ok = m, ng = -1;
  while (llabs(ok - ng) > 1) {
    ll md = (ok + ng) >> 1;
    ll f_md = floor_sum(md+1, m, a, b - c2 + m) - floor_sum(md+1, m, a, b - c1 + m);
    if (f_md > 0) ok = md;
    else ng = md;
  }
  if (ok == m) return INF;
  return a * ok + b;
}

void solve() {
  ll x, y, p, q; cin >> x >> y >> p >> q;
  ll ans = INF;
  chmin(ans, subsolve(p+q, 2*(x+y), x, p+q, p));
  chmin(ans, subsolve(2*(x+y), p+q, p, x+y, x));
  if (ans == INF) cout << "infinity\n";
  else cout << ans << '\n';
}

int main() {
  int t; cin >> t;
  while(t--) solve();
}

余談

$O(\log (X+Y) + \log (P+Q))$ でも解けるようです(すごい)

タイトルとURLをコピーしました