問題概要
いくつかの石が 2 つの山に分けられて積まれている.片方の山に積まれている石の数は $A$ 個で,他方の山のそれは $B$ 個である.
この 2 つの山に対し,次の操作を $K$ 回行う.
- 2 つの山の内,石の個数が多くない方を $X$ とし,他方を $Y$ とする.$Y$ の方の山から $X$ 個の石を他方に移動し,$X$ の方の山の石の数を $2X$ ,$Y$ の方の山の石の個数を $Y - X$ とする
$K$ 回の操作が完了したときの,石の数が多くない方の山に積まれている石の数を求めよ.
$1 \leq A, B \leq 10^9$
$1 \leq K \leq 2 \times 10^9$
解法
$N = A + B$ とします.
$K$ 回の操作を 1 回ずつシミュレーションすると TLE するので,きちんと考察していきます.
1 回の操作について考えてみます.2 つの山の内,大きさを気にしないで片方を選び,そこに積まれている石の個数を $X$ とします.このとき,他方の山に積まれている石の個数は $N - X$ です.操作を 1 回施したときの,$X$ の方の山の変化後の石の個数 $X'$ を考えると,$X$ と $N - X$ の大小関係によって次のように場合分けされます.
- $X \leq N - X$ のとき(i.e. $2X \leq N$ のとき) : $X' = 2X$
- それ以外のとき : $X' = X - ( N - X ) = 2X - N$
この 2 つのケースを見比べてみると,どちらも $X' \equiv 2X \pmod N$ と書けることが分かります.従って,どちらの場合も操作の度に $\mathbb Z /N\mathbb Z $ 上で 2 が掛かるので,$K$ 回操作した後で $A$ の山に積まれている石の個数は $A \times 2^K \bmod N$ となります.
$2^K$ の部分を反復二乗法を用いて先に計算することで,$O( \log K )$ 時間で $K$ 回操作後に片方の山に積まれている石の個数を求めることができます.その個数を $A'$ とすると,答えは(石の個数が多くない方の山に積まれている石の個数なので) $\min( A', N - A' )$ となります.
コード
using LL = long long; // a^x を mod で求める // 反復二乗法 // O( log x ) long long mod_pow( long long a, long long x, long long mod ); class BearPlays { public: int pileSize( int A, int B, int K ) { const int N = A + B; const int res = LL( A ) * mod_pow( 2, K, N ) % N; return min( res, N - res ); } };