[Python] 백준 1517 - 버블 정렬

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

 

문제

N개의 수로 이루어진 수열 A[1], A[2], …, A[N]이 있다. 이 수열에 대해서 버블 소트를 수행할 때, Swap이 총 몇 번 발생하는지 알아내는 프로그램을 작성하시오.

버블 소트는 서로 인접해 있는 두 수를 바꿔가며 정렬하는 방법이다. 예를 들어 수열이 3 2 1 이었다고 하자. 이 경우에는 인접해 있는 3, 2가 바뀌어야 하므로 2 3 1 이 된다. 다음으로는 3, 1이 바뀌어야 하므로 2 1 3 이 된다. 다음에는 2, 1이 바뀌어야 하므로 1 2 3 이 된다. 그러면 더 이상 바꿔야 할 경우가 없으므로 정렬이 완료된다.

입력

첫째 줄에 N(1 ≤ N ≤ 500,000)이 주어진다. 다음 줄에는 N개의 정수로 A[1], A[2], …, A[N]이 주어진다. 각각의 A[i]는 0 ≤ |A[i]| ≤ 1,000,000,000의 범위에 들어있다.

출력

첫째 줄에 Swap 횟수를 출력한다.

---

풀이

대표적인 Inversion Count 문제이다.

Inversion Count란, 어떤 배열이 주어져있을 때, 정렬 순서와 반대되는 원소쌍의 개수를 세는 문제이다.

이번에는 내가 가장 좋아하는 수열인 파이의 자릿수(블로그의 다른 글에도 자주 등장)를 가지고 예시를 들도록 하자

 

[3, 1, 4, 1, 5, 9, 2]

 

결국에는 정렬 순서가 어긋난 쌍이 몇 개 있는지 세는 문제인데,

각 원소에 대해 자신보다 앞에 자기 자신보다 큰 원소가 몇 개 있는지 세면 될 것이다.

위의 배열에서는

3: 0개

1: 1개 (3)

4: 0개

1: 2개 (3, 4)

5: 0개

9: 0개

2: 4개 (3, 4, 5, 9)

 

그래서 정답은 7이 될 것이다.

 

 

N의 범위가 500,000이므로 $N^2$ 풀이는 TLE가 날 것이다.

지난번에 병합 정렬을 이용해서 풀었었는데, 이번에는 세그트리를 공부하고자 세그먼트 트리를 이용해 볼 것이다.

 

 

아이디어는 다음과 같다.

 

세그트리는 구간 합 쿼리와 업데이트 쿼리를 빠르게 처리할 수 있으므로 다음의 과정을 진행한다.

원소의 크기가 큰 원소부터 역으로,

 

1. 0부터 그 원소의 인덱스 바로 앞까지 구간합을 구해서 정답에 더한다

2. 해당 원소의 노드 값을 1 올린다

 

0부터 그 원소의 인덱스 바로 앞까지 구간합을 구하는 것은 현재 자기보다 앞에 있으면서 더 크기가 큰 원소의 개수를 세주는 것이고,

해당 원소의 노드 값을 1 올리는 것은 원소(자기 자신)의 개수를 1 올려주는 것이다.

 

앞에 있으면서 더 크기가 큰 원소의 개수를 세주는 것이기에, 크기가 큰 수부터 쿼리를 진행하는 것이다.

만약 크기가 같은 원소가 둘 이상 있다면 인덱스가 더 큰 것부터 쿼리를 진행한다. 그래야 자신보다 인덱스가 작으면서 크기가 같은 원소들이 세지지 않을 테니까.

 

예를 들어 3, 1, 4, 1, 5, 9, 2에서 퀴리가 진행되는 순서는 아래와 같다.

 

이 인덱스들을 빨리 찾아주기 위해 아래와 같이 정렬을 딱 한 번 하였다.

tmp = [(arr[i], i+1) for i in range(N)]
tmp.sort(reverse=True)

 

코드는 Bottom-Up 세그트리를 이용하였고, 전체 구현 코드는 아래와 같다.

 

N = int(input())
arr = list(map(int, input().split()))

size = 1
while size < N:
    size <<= 1
tree = [0]*(2*size)
size -= 1

def update(idx, diff):
    idx += size
    while idx:
        tree[idx] += diff
        idx //= 2

def search(l, r):
    l += size
    r += size
    ans = 0
    while l <= r:
        if l%2 == 1:
            ans += tree[l]
            l += 1
        l //= 2

        if r%2 == 0:
            ans += tree[r]
            r -= 1
        r //= 2
    return ans


cnt = 0
tmp = [(arr[i], i+1) for i in range(N)]
tmp.sort(reverse=True)
#print(tmp)
for _, i in tmp:
    cnt += search(1, i)
    update(i, 1)
print(cnt)