AtCoder: B - Kleene Inversion - 파이썬(Python)

https://atcoder.jp/contests/jsc2019-qual/tasks/jsc2019_qual_b

B - Kleene Inversion

<문제 선지>

$N$개의 정수 수열 $A = A_0, A_1, ... , A_{N-1}$ 가 있다고 하자.
$B$는 $A$의 $K$개의 사본을 연결하여 만든 $K \times N$ 개의 정수로 이루어진 수열이다.
예를 들어, $A = 1, 3, 2$이고 $K = 2$이면 $B = 1, 3, 2, 1, 3, 2$이다.
$B$의 반전 수에서 $10^9+7$의 나머지를 구하라.
여기서 $B$의 반전 수는 $B_i​>B_j$​인 정수 $(i, j) (0≤i<j≤K×N−1)$의 순서쌍의 수로 정의한다.

<접근 방법>

조합을 이용하여 문제를 해결할 수 있다.

 

$B = a_0a_1a_2...a_{N-1}$ 이라고 하자. 여기서, $a_i = A$를 의미한다.

그러면 가능한 반전 수 유형은 총 2가지로 분류할 수 있다.

 

• $a$ 내부에서의 반전 수

$a$ 내부에서의 반전 수는 총 $N$개 존재한다. 왜냐하면 $B = a_0a_1a_2...a_{N-1}$에서 $a$의 개수는 $N$개이기 때문이다.

• 서로 다른 $a$ 간에서의 반전 수

서로 다른 $a$ 간에서의 반전 수는 총 ${}_{N}C_{2} = \frac{N(N - 1)}{2}$개 존재한다. 왜냐하면 순서 상관 없이 N개 중, 2개를 선택하는 경우의 수와 같기 때문이다.

 

결론적으로, $(a 내부에서의 반전 수의 합) \times N + (서로 다른 a 간에서의 반전 수) \times \frac{N(N - 1)}{2}$를  $10^9+7$로 나눈 나머지를 구하면 된다.

<소스 코드>

위의 접근법을 구현한 전체 소스 코드이다.

a 자신의 반전 수와 서로 다른 a 간의 반전 수 계산 방식이 다름에 주의해야 한다.

n, k = map(int, input().split())
*arr, = map(int, input().split())
r = 10 ** 9 + 7

res = [0] * n
count = [0] * n
doub_count = [0] * n

# a 자신의 반전 수
for i in range(n):
    for j in range(i + 1, n):
        if arr[i] > arr[j]:
            count[i] += 1

# 서로 다른 a 간의 반전 수
for i in range(n):
    for j in range(n):
        if arr[i] > arr[j]:
            doub_count[i] += 1

print((sum(count) * k + sum(doub_count) * k * (k - 1) // 2) % r)