문제 소개
크기가 n인 배열 A와 두 정수 p, q가 주어져 있다고 가정해 보겠습니다. 우리는 배열 A에 대해 다음 두 가지 연산을 수행할 수 있습니다.
- l < r을 만족하는 두 인덱스 (l, r)를 무작위로 선택한 뒤 A[l]과 A[r]의 값을 서로 교환(swap)합니다.
- l < r을 만족하는 두 인덱스 (l, r)를 무작위로 선택한 뒤 인덱스 l부터 r까지의 부분 배열을 뒤집습니다(reverse).
첫 번째 연산을 p회, 두 번째 연산을 q회 수행한 후, l < r을 만족하는 두 인덱스 l과 r을 다시 무작위로 선택하여 S = 부분 배열 A[l..r]에 포함된 모든 원소의 합이라고 정의합니다. 이때 우리가 구해야 하는 것은 바로 S의 기댓값입니다.
예시 살펴보기
예를 들어 입력이 A = [1, 2, 3], p = 1, q = 1이라면 결과값은 4.667이 됩니다. 그 과정을 단계별로 살펴보겠습니다.
1단계: 스왑 연산에는 세 가지 경우의 수가 있습니다.
- swap(0, 1) → 배열은 [2, 1, 3]
- swap(0, 2) → 배열은 [3, 2, 1]
- swap(1, 2) → 배열은 [1, 3, 2]
2단계: 각 결과마다 뒤집기 연산 역시 세 가지 경우의 수가 있습니다.
- [2, 1, 3] → [1, 2, 3], [3, 1, 2], [2, 3, 1]
- [3, 2, 1] → [2, 3, 1], [1, 2, 3], [3, 1, 2]
- [1, 3, 2] → [3, 1, 2], [2, 3, 1], [1, 2, 3]
가능한 배열은 총 9가지이므로 각 배열이 나타날 확률은 1/9입니다. 또한 각 배열마다 동일한 확률을 가지는 3가지 부분 배열의 합이 존재합니다. 예를 들어 [1, 2, 3]의 경우 1+2, 2+3, 1+2+3의 세 가지 합을 얻을 수 있습니다. 따라서 이 입력에서 가능한 전체 결과는 27가지이며, 기댓값은 모든 27개의 S 값 총합을 27로 나누어 계산할 수 있습니다.
접근 방법
이 문제는 각 위치별 원소 분포를 확률 행렬로 표현한 뒤 행렬 곱셈을 반복적으로 적용하면 효율적으로 해결할 수 있습니다. 먼저 행렬과 벡터의 곱을 계산하는 함수를 정의합니다.
- 행렬 a, 벡터 v, 크기 n을 매개변수로 받는 matmul() 함수를 정의합니다.
- toret := 크기가 n이고 0으로 초기화된 배열로 설정합니다.
- i를 0부터 n-1까지 순회하며, 각 i에 대해 j를 0부터 n-1까지 순회하면서 toret[i] := toret[i] + a[i][j] * v[j]를 누적합니다.
- toret을 반환합니다.
메인 solve() 함수에서는 다음 단계를 수행합니다.
- n := 배열 A의 크기
- temp := 새로운 리스트
- swp := (n - 3) / (n - 1)
- swapvalp := ((swp^p) * (n - 1) + 1) / n
- swapvalm := (1 - swp^p) / n
- rev := 비어 있는 새 리스트
- dotv := 비어 있는 새 리스트
- i를 0부터 n-1까지 순회하며 다음을 반복합니다.
- swaprow := 새 빈 리스트, revrow := 새 빈 리스트로 생성
- j를 0부터 n-1까지 순회하며 swaprow의 끝에 swapvalm을, revrow의 끝에 2 * (min(i, j, n-i-1, n-j-1) + 1) / (n * (n - 1))을 추가합니다.
- swaprow[i] := swapvalp로 설정
- revrow[i] := 1.0 - 2 * ((i + 1) * (n - i) - min(i + 1, n - i)) / (n * (n - 1))로 설정
- temp의 끝에 swaprow를, rev의 끝에 revrow를 추가합니다.
- dotv의 끝에 2 * ((i + 1) * (n - i) - 1) / (n * (n - 1))을 추가합니다.
- A := matmul(temp, A, n)으로 스왑 연산의 효과를 한 번에 반영합니다.
- q회 반복하며 A := matmul(rev, A, n)을 수행해 뒤집기 연산을 적용합니다.
- tot := 0.0으로 초기화한 뒤, i를 0부터 n-1까지 순회하며 tot := tot + dotv[i] * A[i]를 누적합니다.
- tot을 최종 결과로 반환합니다.
구현 예제
아래의 Python 구현 예제를 통해 더욱 명확하게 이해해 보겠습니다.
def matmul(a, v, n):
toret = [0]*n
for i in range(n):
for j in range(n):
toret[i] += a[i][j]*v[j]
return toret
def solve(A, p, q):
n = len(A)
temp = []
swp = (n - 3)/(n - 1)
swapvalp = (pow(swp, p)*(n - 1) + 1)/n
swapvalm = (1 - pow(swp, p))/n
rev = []
dotv = []
for i in range(n):
swaprow = []
revrow = []
for j in range(n):
swaprow.append(swapvalm)
revrow.append(2*(min(i, j, n - i - 1, n - j - 1) + 1)/(n*(n - 1)))
swaprow[i] = swapvalp
revrow[i] = 1.0 - 2*((i + 1)*(n - i) - min(i + 1, n - i))/(n*(n - 1))
temp.append(swaprow)
rev.append(revrow)
dotv.append(2*((i + 1)*(n - i) - 1)/(n*(n - 1)))
A = matmul(temp, A, n)
for _ in range(q):
A = matmul(rev, A, n)
tot = 0.0
for i in range(n):
tot += dotv[i]*A[i]
return tot
A = [1,2,3]
p = 1
q = 1
print(solve(A, p, q))
입력
[1, 2, 3], 1, 1
출력
4.667
이처럼 확률 행렬과 행렬 곱셈을 활용하면 모든 경우의 수를 일일이 나열하지 않고도 부분 배열 합의 기댓값을 효율적으로 계산할 수 있습니다.