[백준] 18194 - Bad Hair Day와 기댓값

[문제]

https://www.acmicpc.net/problem/18194

[풀이]

소들이 일렬로 설 수 있는 모든 경우를 확인하면 오래 걸릴 것이다.
문제에서 요구하는 값은 기댓값이기 때문에 조합론으로 접근할 수 있다.

N 마리의 소 중 어떤 소 하나를 선택했다고 하자.
그 소보다 키가 크거나 같은 소의 머리는 확인할 수 없다.
그 소보다 키가 작은 소가 k 마리 있다고 하자.
그러면 그 소보다 키가 크거나 같은 소는 N-k 마리다.
N-K 마리의 소를 적당히 일렬로 세우고, 그 사이에 k 마리의 소를 잘 넣는다고 하면
그 소가 볼 수 있는 소의 마릿수의 기댓값은 $\frac{k}{N-k+1}$ 이다.

모든 소에 대한 기댓값을 모두 더해주면 최종적으로 요구되는 $\frac{P}{Q}$ 값을 구할 수 있다.

문제를 맞추기 위해서는 일차 합동식 $QX \equiv P$ $(mod$ $10^9+7)$ 을 만족하는 $X$ 를 찾아야 한다.
풀이를 쉽게 하기 위해 $10^9+7$을 $C$ 라 하자. 또한, $C$ 는 소수이다.

위의 일차 합동식의 해가 존재하려면 $gcd(Q,C) | P$ 를 만족해야 한다.
문제에서 $P$, $Q$, $X$ 는 반드시 존재하고, $X$ 가 유일하다고 했기 때문에,
$gcd(Q,C)=1$ 라고 생각하고 선형 방정식 $Qx+Cy=gcd(Q,C)=1$ 을 풀어야 한다.
$x$를 구하고 $P$배를 하면 원하는 $X$를 얻을 수 있다.
선형 방정식을 풀기 위해 확장된 유클리드 알고리즘을 사용해도 되지만,
시간 복잡도를 낮추기 위해 페르마의 소정리 및 분할 정복을 통한 거듭제곱으로
$x=Q^{P-2}\%C$ 임을 알 수 있다. 즉, $X=Q^{P-2} \cdot P\%C$ 이다.

[코드]

[Python3]

import sys;r=sys.stdin.readline

def inv(n):
    p=n; res=1; y=C-2
    while y:
        if y%2: res=res*p%C
        p=p*p%C
        y//=2
    return res

C = 10**9 + 7
N=int(r())
cow=list(map(int,r().split()))
cow.sort()

tmp=cow[0]
ans, k = 0, 0
for i in range(1,N):
    if cow[i] != tmp:
        k=i
        tmp=cow[i]
    ans+=inv(N-k+1)*k%C
    ans%=C
print(ans)

어떤 소보다 작은 소의 마릿수인 k를 구하기 위해 정렬을 했다.