[백준 2517] 2012 정올 고등부 2차 "달리기"

문제 출처 : https://www.acmicpc.net/problem/2517

2517번: 달리기

이 문제는 2012년 정보 올림피아드 고등부 문제 입니다.

 

앞에서부터 자신보다 실력이 높은 사람의 수를 세어 자신의 최고 등수를 찾는 문제 입니다. 단순한 문제이지만 입력의 숫자가 많아 문제를 풀기는 단순하지 않습니다.

등수 찾기

맨 첫 번째 선수는 무조건 1등을 할 수 있습니다. 그리고 다음 선수는 첫 번째 선수보다 실력이 높다면 1등, 실력이 낮다면 2등 입니다. 즉 어떤 선수의 차례가 되었을 때 자신보다 실력이 높은 선수의 합을 구하면 자신의 순서를 알 수 있게 됩니다. 예제를 보면서 이해해 보겠습니다.

2, 8, 10, 7, 1, 9, 4, 15의 실력을 가진 선수들이 순서대로 들어 옵니다.

가장 먼저 2의 능력을 가진 선수가 들어옵니다. 자신보다 빠른 선수가 없기 때문에 무조건 1등을 할 수 있습니다. 다음으로 8의 능력을 가진 선수가 들어옵니다. 

8보다 큰 실력을 가진 선수가 없기 때문에 2번 선수도 1등을 할 수 있습니다. 다음으로 10의 능력을 가진 3번 선수가 들어옵니다.

10의 능력의 선수 역시 자신보다 빠른 선수가 없기 때문에 1등을 할 수 있습니다. 다음으로 7의 능력을 가진 선수가 들어옵니다.

자신을 포함하여 자신보다 빠른 선수가 3명 있습니다. 8과 10의 능력을 가진 선수는 이길 수 없기 때문에 최고 등수는 자신보다 큰 숫자들의 합인 3등이 됩니다.

1의 능력을 가진 선수가 들어오면 앞서 나왔던 모든 선수를 이길 수 없습니다. 최고 등수는 자신보다 큰 숫자들의 합 5와 같습니다.

9의 능력을 가진 선수가 들어오면 자신보다 큰 숫자들의 합인 2등이 될 수 있습니다.

4의 능력을 가진 선수의 최고 등수는 4보다 큰 숫자들의 합계인 5등이 됩니다.

마지막으로 15의 능력을 가진 선수가 들어오면 합계가 1이기 때문에 1등을 할 수 있습니다. 이렇게 업데이트가 되면서 누적합계를 구하기 위해서는 세그먼트 트리를 사용하면 문제를 해결하면 됩니다.

좌표 압축하기

문제를 해결할 방법은 찾았지만 한 가지 문제를 더 해결해야 합니다. 바로 실력이 1 이상 1,000,000,000 이하라는 것입니다. 무려 10억이라는 배열을 세그먼트 트리를 만들기에는 너무 숫자가 큽니다. 다행히 선수의 수는 10억보다는 훨씬 적은 50만 입니다. 

좌표의 개수가 50만 이라면 세그먼트 트리를 만드는 것도 가능하다는 생각을 할 수 있습니다. 다음으로 고려해야 할 점은 선수의 실력의 숫자들이 의미가 있느냐는 것입니다. 이 문제에서는 선수의 실력은 중요하지 않습니다. 중요한 것은 몇등인지만 알고 있으면 됩니다. 따라서 좌표압축을 통해 문제를 해결할 수 있습니다. 다빈치코딩 알고리즘에 좌표 압축을 설명해 놓았습니다. 좌표 압축이 무엇인지 잘 모른다면 해당 링크를 확인해 보시기 바랍니다. 

알고리즘

위에서 생각한 방식의 알고리즘은 좌표 압축을 한 후 세그먼트 트리를 이용하면 될 것 같습니다. 세그먼트 트리의 누적합 구하기, 누적합 업데이트를 통해 생각할 수 있습니다. 여기서 발생하는 문제는 선수의 실력이 1부터 1,000,000,000 까지 입니다. 이 많은 숫자들을 세그먼트 트리로 만들기는 어렵습니다. 다행히 선수의 수는 50만 이하이고, 모든 선수의 실력은 다르다고 합니다. 따라서 이것을 이용하여 좌표압축을 하여 세그먼트 트리를 만들 수 있습니다.

1. 선수들의 정보를 입력 순서와 실력으로 정렬합니다.
2. 정렬된 실력을 좌표 압축으로 줄여 줍니다.
3. 입력 순서로 다시 정렬 합니다.
4. 첫 번째 선수부터 선수들의 실력을 세그먼트 트리에 등록합니다.
5. 자신을 포함한 자신보다 실력이 높은 선수들의 합을 구합니다.
6. 5번의 합이 곧 순위 입니다. 해당 순위를 출력합니다.
7. 4번으로 돌아가 다음 선수를 확인 합니다.

그럼 이 순서대로 코드를 작성해 보겠습니다.

 

입력 받기

N = int(input())

arr = []

for i in range(N):
    arr.append([i, int(input())])

입력 받을 선수의 수 N을 입력 받습니다. 다음으로 선수의 실력을 입력 받습니다. 이 때 어떤 순서로 입력되었는지가 중요하기 때문에 인덱스값 i와 함께 저장합니다.

좌표 압축하기

arr.sort(key=lambda x : x[1])

for i in range(N):
    arr[i][1] = i

arr.sort()

먼저 실력을 기준으로 정렬해야 합니다. lambda 함수를 통해 실력인 x[1]을 기준으로 정렬 합니다. lambda에 대해 잘 모른다면 다빈치코딩 파이썬 lambda 함수를 확인해 보시기 바랍니다.

정렬을 하였기 때문에 0부터 N-1까지의 숫자로 선수의 실력을 바꿔줍니다. 그리고 다시 인덱스값을 기준으로 정렬합니다. 2차 배열 이상을 정렬할 때 정렬 기준을 입력하지 않으면 첫 번째 배열을 기준으로 정렬이 됩니다. 인덱스 값으로 정렬하면 되기 때문에 따로 key를 지정하지 않아도 정렬이 됩니다.

 

바텀업 세그먼트 트리

세그먼트 트리는 바텀업 세그먼트 트리로 만들 것입니다. 파이썬은 느리기 때문에 가능하면 바텀업 세그먼트 트리로 만드는 것이 문제를 해결하는 데 도움이 됩니다. 다빈치코딩 알고리즘에 바텀업 세그먼트 트리에 대해 설명해 놓았습니다. 잘 모른다면 설명 먼저 확인해 보시기 바랍니다.

 

트리 초기화

size = 1
while size <= N:
    size <<= 1

tree = [0] * (2 * size)

세그먼트 트리를 만들기 위해 트리의 크기를 구하고, 크기를 바탕으로 트리를 초기화 해주었습니다.

 

구간 합계 함수

def query(num):
    left = num + size
    right = N - 1 + size

    ans = 0
    while left <= right:
        if left % 2 == 1:
            ans += tree[left]
            left += 1
        left //= 2

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

입력된 수 num부터 끝까지 구간합을 구하는 함수 입니다. 함수 내용은 바텀업 세그먼트 트리와 같습니다.

업데이트 함수

def update(i):
    idx = i + size

    tree[idx] = 1
    idx //= 2
    while idx:
        tree[idx] = tree[2 * idx] + tree[2 * idx + 1]
        idx //= 2

선수의 실력이 입력될 때마다 선수의 실력으로 누적합을 만드는 함수 입니다. 이것 역시 바텀업 세그먼트 트리에 설명한 내용과 같습니다.

출력하기

for i in range(N):
    num = arr[i][1]
    update(num)
    print(query(num))

선수의 실력이 입력되면 해당 선수의 실력을 인덱스로 업데이트 해줍니다. 그리고 입력된 숫자부터 끝까지 구간의 합을 구하여 출력하면 그것이 등수가 됩니다.

전체 코드

전체 코드를 확인해 보겠습니다. 

N = int(input())

arr = []

for i in range(N):
    arr.append([i, int(input())])

arr.sort(key=lambda x : x[1])

for i in range(N):
    arr[i][1] = i

arr.sort()

size = 1
while size <= N:
    size <<= 1

tree = [0] * (2 * size)

def update(i):
    idx = i + size

    tree[idx] = 1
    idx //= 2
    while idx:
        tree[idx] = tree[2 * idx] + tree[2 * idx + 1]
        idx //= 2

def query(num):
    left = num + size
    right = N - 1 + size

    ans = 0
    while left <= right:
        if left % 2 == 1:
            ans += tree[left]
            left += 1
        left //= 2

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

for i in range(N):
    num = arr[i][1]
    update(num)
    print(query(num))