백준 1777번 순열복원 - 스위프트(Swift) 풀이 + 그림 설명

 1~N에 대해 i보다 뒤에 나오면서 더 작은 수의 개수가 주어졌다. 따라서 큰 수부터 내림차순으로 뒤에 빈자리가 inversionSequence[i]개 있게 되는 곳에 놓으면 된다. 큰 수부터 고려해주고 있기 때문에, 빈자리의 의미는 앞으로 더 작은 수가 위치할 자리이기 때문이다.

 

 

 먼저 가장 큰 5부터 자리를 찾아주자. 5의 Inversion Count가 2이므로 뒤에 빈자리가 2개가 되는 곳에 놔야 한다.  그림으로 보면 3번째 자리가 되겠다.

 

 

 그럼 뒤에 빈자리가 x개가 되는 곳을 어떻게 찾을까? 순차 탐색으로 뒤에서부터 하나씩 세면 당연히 O(N^2)으로 시간 초과가 난다. 세그먼트 트리로 시간문제를 해결할 수 있다. 해당 자리가 비어있는지 여부를 배열로 둬서, 맨 뒤에서부터의 index까지의 구간합이 최초로 x+1 이상이 되는 위치를 찾으면 된다. (index에 놓으면, 이제 뒤에 빈자리가 x개 있을 것이기 때문에.) 이 lower bound 탐색에 O(logN)이 들고, 구간합을 구하는데 O(logN)이 들기 때문에 최종 시간 복잡도는 O(N(logN)^2)이 된다!

 

 

 이제 4 차례. 3번째 자리는 방금 5를 놨기 때문에 놨다고 0으로 체크해준다. 4의 Inversion Count가 1이므로 뒤에서부터의 구간합이 최초로 2가 되는 곳을 찾아야 한다.

 

 

 다음 3 차례. Inversion Count가 2이므로 뒤에서부터의 구간합이 3이 되는 위치에 놓아야 한다.

 

 

 같은 방식으로 1까지 채워주면, [3, 2, 5, 4, 1]이라는 순열로 복원된다. 각각 inversion counting을 구해보면 [0, 1, 2, 1, 2]가 나오는 것을 확인할 수 있다.

 

import Foundation

class FenwickTree {
    var n: Int
    var tree = Array(repeating: 0, count: 100002)
    
    init(n: Int) {
        self.n = n
        (1...100000).forEach{ add(at: $0, 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)
        }
    }
    
    // [1...]부터의 구간합이 최초로 targetSum 이상이 되는 위치의 인덱스 반환
    func lb(targetSum: Int) -> Int {
        var lo = 1, hi = n + 1
        
        while lo < hi {
            let mid = (lo + hi) / 2
            if sum(to: mid) < targetSum { lo = mid + 1 }
            else { hi = mid }
        }
        
        return lo
    }
}

let N = Int(readLine()!)!
var arr = readLine()!.split(separator: " ").map{ Int(String($0))! }

var fwt = FenwickTree(n: N)
var seq = Array(repeating: 0, count: N + 1)

for i in 0..<N {
    let num = N - i                                     // 큰 수부터 차례로 위치시키기
    let index = fwt.lb(targetSum: arr[num - 1] + 1)     // num보다 뒤에 있으면서 작은 것의 개수가 arr[num - 1]개인 자리
    seq[N - index] = num                                // 뒤에서 index번째에 위치하면 됨
    
    fwt.add(at: index, val: -1)                         // 뒤에서 index번째 자리 선점 여부 체크
}

var result = ""
for i in 0..<N {
    result.write("\(seq[i]) ")
}
print(result)