백준 2517번 달리기 - 스위프트(Swift) 풀이 + 그림 설명

 문제를 요약하면, 각 선수에 대해 나보다 앞에 있으면서 큰 사람의 수 + 1을 구해라 이다. N이 500,000이므로 브루트포스로 2중 for문을 도는 O(N^2)으로는 당연히 시간 초과가 난다. 이런 Inversion Counting 문제는 세그먼트 트리(펜윅트리)로 해결 가능하다. (머지 소트를 쓰는 방법도 있다.)

 

 아래와 같은 상황이라고 하자. 나보다 앞에 있으면서 큰 사람의 수를 카운트해주기 위해서, 작은 수부터 오름차순으로 고려해줄 것이다.

 

 

 그리고 아직 고려 안했는지 여부를 저장하는 배열을 하나 더 두자. 아직 고려를 안 했으면 1이고 고려했으면 0이다. 따라서 초기에는 전부 1.  이제 실력이 낮은 사람부터 [1, 7, 9, 23, 56] 순으로 아직 고려 안했는지 여부[1...(index-1)]의 구간합 + 1을 구해주면 된다. 아직 고려 안 했다는 것은 나보다 실력이 높다는 것이기 때문에! 실제로 해보면..

 

 

 먼저 가장 낮은 1부터 찾자. 1의 index는 1이고, 앞에 아무것도 없으므로 구간합 0이다. 따라서 최대 0 + 1 = 1등을 할 수 있다.

 

 

 다음은 7인데 넘어가기 전에 방금 고려한 1을 고려했다고 체크해주자. 실력이 7인 사람의 index는 3이고, 앞에 있는 [1...2]의 구간합을 구하면 1이다. 즉, 7보다 앞에 있으면서 7보다 큰 사람은 한 명뿐이므로 최대 1 + 1 = 2등을 할 수 있다.

 

 

 9까지만 해보면, 실력이 9인 사람의 index는 5이고, 앞에 있는 [1...4]의 구간합을 구하면 2이다. 즉, 9보다 앞에 있으면서 큰 사람은 2명이므로, 최대 2 + 1 = 3등을 할 수 있다.

 

 

 실력이 낮은 사람부터 한 명씩 인덱스를 찾아내는 방법은, 처음에 (인덱스, 실력)을 쌍으로 묶어서 정렬을 해주면 된다.

 

import Foundation

class FenwickTree {
    var tree = Array(repeating: 0, count: 500001)
    
    init() {
        for i in 1...500000 { self.add(at: i, val: 1) }
    }
    
    func sum(to: Int) -> Int {
        var ret = 0
        var pos = to
        
        while (pos > 0) {
            ret += tree[pos]
            pos &= (pos - 1)
        }
        
        return ret
    }
    
    func add(at: Int, val: Int) {
        var pos = at
        
        while(pos < tree.count) {
            tree[pos] += val
            pos += (pos & -pos)
        }
    }
}

var arr = Array(repeating: (1000000001, 1000000001), count: 500001)
var rank = Array(repeating: -1, count: 500001)

let N = Int(readLine()!)!
for i in 0..<N { arr[i] = (i + 1, Int(readLine()!)!) }

let sortedArr = arr.sorted(by: { $0.1 < $1.1 })

let fwt = FenwickTree()
for i in 0..<N {
    let index = sortedArr[i].0
    rank[index] = fwt.sum(to: index - 1) + 1    // 최선의 등수 = i보다 앞에 있으면서 i보다 큰 것의 개수 + 1
    fwt.add(at: index, val: -1)
}

var result = ""
for i in 1...N {
    result.write("\(rank[i])\n")
}
print(result)