BOJ 18936 신기한 공놀이
문제 설명
$A$개의 공 중 $B$개의 빨간색 공이 있는데, 무작위로 두 개의 공을 뽑았을 때 적어도 하나는 빨간색이 아닐 확률이 $1 \over n^{2}$이 되는 $(A, B)$ 쌍 중 $A$가 $m$번째로 작은 쌍을 구하는 문제이다.
식 변형
문제의 조건은 다음과 같다. \[1-{b(b-1) \over a(a-1)} = {1 \over n^{2}}\]
양변에 $n^{2}a(a-1)$을 곱해 정리하면 다음과 같다. \[(n^{2}-1)a(a-1) = n^{2}b(b-1)\]
일차항을 소거하기 위해 완전제곱식을 만들어주면 다음과 같다. \[(n^{2}-1)({a^2}-a+{1 \over 4})-{ {n^{2}-1} \over 4 } = n^{2}({b^2}-b+{1 \over 4})-{n^{2} \over 4}\]
양변에 4를 곱하고, $x=n(2b-1), y=2a-1, d=n^{2}-1$로 치환하면 다음과 같다. \[{x^2}-d{y^2}=1\] \[x \equiv n (mod \space 2n)\] \[y \equiv 1 (mod \space 2)\]
우리는 이 식을 푸는 것이 목표이다. 일반적으로 제곱수가 아닌 자연수 $d$에 대해 \({x^2}-d{y^2}=1\) 형태의 정수 방정식을 펠 방정식(Pell’s equation)이라 부른다.
펠 방정식
펠 방정식을 푸는 법을 알아보자. ${x^2}-d{y^2}=1$인 모든 자연수해를 찾는 것이 목표이다.
해 $(x_1, y_1)$을 하나 찾았다고 하자. $\sqrt{d}$가 무리수일 때, 다음 식이 성립한다. \[(x_{1}+\sqrt{d}y_{1})^{n} = x_{n}+\sqrt{d}y_{n}\] \[(x_{1}-\sqrt{d}y_{1})^{n} = x_{n}-\sqrt{d}y_{n}\]
두 식을 변변 곱하면, ${x_{1}}^{2}-d{y_{1}}^{2} = 1$이므로 ${x_{n}}^{2}-d{y_{n}}^{2} = 1$이 되어 $(x_{n}, y_{n})$ 역시 해가 된다.
$(x_{1}, y_{1})$이 $x$가 가장 작은 해라면, 위의 방법으로 생성되는 $(x_{n}, y_{n})$이 이 방정식의 유일한 해가 된다.
저 쌍이 방정식의 해가 됨은 자명하지만, 저 쌍 외의 해가 없음을 보이는 것은 자명하지 않다. 증명은 다음과 같다.
$(x, y)$가 주어진 방정식의 해라고 하자. $\displaystyle \lim_{n -> \infty} (x_{1}+\sqrt{d}y_{1})^{n} = \infty$이므로, $x_{n}+\sqrt{d}y_{n} \leq x+\sqrt{d}y < x_{n+1}+\sqrt{d}y_{n+1}$인 유일한 자연수 $n$이 존재한다. 따라서 \[1 \leq {x+\sqrt{d}y} \over {x_{n}+\sqrt{d}y_{n}} < {x_{n+1}+\sqrt{d}y_{n+1}} \over {x_{n}+\sqrt{d}y_{n}} = x_{1}+\sqrt{d}y_{1}\]
${ {x+\sqrt{d}y} \over {x_{n}+\sqrt{d}y_{n}} } = (x+\sqrt{d}y)(x_{n}-\sqrt{d}y_{n}) = a+b\sqrt{d}$라 하면, $(a, b)$ 역시 주어진 방정식의 해가 된다.(전개해보라) 이제 다음을 관찰해보자.
- $(a+b\sqrt{d})(a-b\sqrt{d}) = 1$
- $a+b\sqrt{d} \geq 1 > 0$
- $a+b\sqrt{d} < x_{1}+y_{1}\sqrt{d}$
1, 2로부터 $a-b\sqrt{d} > 0$이고, $2a = a+b\sqrt{d} + a-b\sqrt{d} >0$이다. 따라서 $a>0$이다.
1, 3으로부터 $a-b\sqrt{d} \leq a+b\sqrt{d}$이고, $b \geq 0$이다.
$b>0$이면, $(a, b)$ 또한 주어진 방정식의 해인데 $(x_{1}, y_{1})$이 최소해임에 모순이다. 따라서 $b=0, a=1$이고, 이는 $(x, y)=(x_{n}, y_{n})$임을 의미한다. 따라서 모든 해는 위와 같은 꼴이다.
문제로 돌아가서
따라서 주어진 문제를 풀기 위해서는 우선 최소해를 찾아야 한다. 다행히 $d=n^{2}-1$이므로, $(n, 1)$이 최소해임을 쉽게 알 수 있다.
한 가지 조건이 더 있다. $a, b$가 자연수여야 하므로, $x \equiv n (mod \space 2n), y \equiv 1 (mod \space 2)$을 만족해야 한다.
$mod \space 2n$ 상에서 첫 $(x_i, y_i)$를 관찰해보면 다음과 같다. \[(x_{1}, y_{1}) \equiv (n, 1)\] \[(x_{2}, y_{2}) \equiv (-1, 0)\] \[(x_{3}, y_{3}) \equiv (n, -1)\] \[(x_{4}, y_{4}) \equiv (1, 0)\] \[(x_{5}, y_{5}) \equiv (n, 1)\]
주기 4로 반복됨을 알 수 있고, 그중 $x \equiv n (mod \space 2n), y \equiv 1 (mod \space 2)$인 경우는 $(x_{2k-1}, y_{2k-1})$인 경우이다.
우리는 $m$번째 해를 구하는 것이 목표이므로, $k=m$으로 택하면 해를 구할 수 있다. 거듭제곱은 $O(log m)$ 시간에 계산이 가능하므로, 이 문제는 $O(Qlogm)$ 시간에 풀 수 있다.
$n=1$일 때, $d=0$인데, 따라서 펠 방정식의 형태는 아니지만 $x^2 = 1$이라는 자명한 방정식이다. 이 경우도 최소해를 $(1, 1)$이라 두고 위의 과정을 똑같이 따라하면 해를 구할 수 있다. 구현은 아주 간단하다.
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> pii;
ll n, m, k, ans;
const ll mod=1e9+7;
ll d;
inline pii mul(pii a, pii b)
{
return pii((a.first*b.first+d*a.second%mod*b.second)%mod, (a.first*b.second+a.second*b.first)%mod);
}
ll pw(ll a, ll b)
{
ll ret=1;
while(b)
{
if(b&1) ret=ret*a%mod;
a=a*a%mod;
b>>=1;
}
return ret;
}
int main()
{
ios_base::sync_with_stdio(false);
cin.tie(0);
ll i, j, temp=0;
ll q;
cin>>q;
while(q--)
{
cin>>n>>m;
d=(n*n-1)%mod;
pii r(1, 0), a(n, 1);
m=2*m+1;
while(m)
{
if(m&1) r=mul(r, a);
a=mul(a, a);
m>>=1;
}
cout<<(r.second+1)*pw(2, mod-2)%mod<<' '<<(r.first+n)*pw(2*n%mod, mod-2)%mod<<'\n';
}
}
주절주절
이 문제는 학교 기초통계학 시간에 선생님께서 주신 문제를 일반화한 것이다. 그땐 아마 확률이 $1/2$이었나 그랬던 것 같다. 분모가 $1 \over k$이면, $kx^2-(k-1)y^2=1$꼴의 방정식을 풀어야 하는데, 나는 이걸 풀 줄 모른다. 대신 $k, k-1$중 하나가 제곱수인 경우는 위에서 언급한 펠 방정식이 되어 쉽게 풀 수 있는 형태가 된다. 그래서 $k=n^{2}$로 문제를 냈고, $k=n^{2}+1$이어면 $x^{2} - dy^{2} = -1$이 되는데, 거의 똑같이 풀린다.
이렇게 만든 문제를 2019 Scioi에 출제했고, 많이 풀리진 않았는데 열심히 작은 경우의 답을 구해 규칙을 찾아 문제를 푼 사람이 등장해서 놀랐다.
백준 스택 권한을 얻고 이 문제를 가장 먼저 출제하였다. 검수자로는 대회 당시 AC를 받은 사람들을 섭외했다.
내 코드는 148ms에 돌고, 검수자들 중 가장 빠른 코드는 124ms에 돌았다. $x_{i}, y_{i}$는 전 항의 점화식으로 표현되므로 간단히 $2 \times 2$ 행렬로 나타내 행렬 곱셈 거듭제곱으로 풀 수 있는데, 이러면 $(n+\sqrt{n^{2}-1})$을 거듭제곱하는 것 보다 상수가 커지게 된다. (각각 대충 8, 4 정도일 것이다) 원래는 이 풀이가 TLE가 났으면 좋겠다고 생각했는데, 상수 2배 차이라 막을 별다른 방법 없이 그냥 500ms 제한으로 출제했다.
그랬더니 크기 3짜리 점화식을 찾아서 TLE를 받다가 FastIO 등으로 열심히 커팅해 통과하는 분들이 등장했고, 막자고 TL을 300ms 같이 하는건 좀 추한것 같아서 그냥 냅뒀다. 그러다 보니 펠 방정식을 풀 줄 몰라도 이 문제 자체는 쉽게 풀 수 있는 것 같다.