백준 1572번 중앙값 - 스위프트(Swift) 시간초과 해결

 스위프트로 문제 풀면 절반은 TLE를 맞는.. 오늘 벌써 두 번째 TLE인데 이번엔 좀 신기한 상황이었다. 이 문제 바로 전에 1615번 문제에서 이유를 알 수 없는 TLE를 맞았는데, Swift로 맞은 사람이 없길래 포기하고 바로 C++로 풀었는데 이 문제는 Swift로 맞은 사람이 있었다...! 결국 언어 한계가 아니라 내 코드 문제라는 건데.. 코드가 길지도 않은데 대체 어디서 시간을 잡아먹는 건지 알 수가 없었다.

 예전에 C++로 풀이할 때, 메모리 접근 시간 때문에 TLE를 맞은 적이 있었다. 로직은 그대로 두고 2차원 배열 행과 열만 바꿔서 저장했더니 통과가 됐었다. 그래서 설마 이것도 메모리 접근에서 시간을 잡아먹는 건가 싶어 클래스 안에 타고 들어가는 시간 줄이려고 클래스로 따로 뺐던 펜윅트리 내부의 tree배열을 전역으로 빼줬다. 또 arr랑 tree랑 번갈아가면서 자주 접근하는데, 둘이 같이 전역으로 선언해줘야 둘이 가까운 곳에 할당되지 않을까 싶어서...

 결과는......통과가 됐다! ㄴㅇㄱ PS 할 때는 객체지향은 머리에서 지워야겠다:) 왠지 아까 Swift로 포기한 문제도 이렇게 바꾸면 될 것 같다.

 

 시간초과 코드

import Foundation

class FenwickTree {
    var n: Int
    var tree = Array(repeating: 0, count: 65538)
    
    init() { self.n = 65537 }
    
    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 input = readLine()!.split(separator: " ").map{ Int(String($0))! }
let N = input[0]
let K = input[1]
let mid = (K + 1) / 2

var arr = Array(repeating: 1, count: N + 1)
for i in 1...N { arr[i] = Int(readLine()!)! + 1 }       // 0도 들어올 수 있는데, 펜윅트리에서 index는 양수여야 하므로 +1을 해서 기록해줌

let fwt = FenwickTree()
for i in 0...K-1 {                                      // 일반화를 위해 맨 앞에 0이 추가로 있다고 가정
    fwt.add(at: arr[i], val: 1)
}

var result: Int64 = 0
for i in K...N {
    fwt.add(at: arr[i - K], val: -1)                    // arr[i - K]는 이번 부분 수열 범위에서 제외됨
    fwt.add(at: arr[i], val: 1)                         // arr[i]는 이번 부분 수열 범위에 추가됨
    result += Int64(fwt.lb(targetSum: mid) - 1) // 작은 수 부터 카운트했을 때 합이 (K+1)/2개 이상이 되는 수가 중앙값
}
print(result)

 

 통과된 코드

import Foundation

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 = 65538
    
    while lo < hi {
        let mid = (lo + hi) / 2
        if sum(to: mid) < targetSum { lo = mid + 1 }
        else { hi = mid }
    }
    
    return lo
}

let input = readLine()!.split(separator: " ").map{ Int(String($0))! }
let N = input[0]
let K = input[1]
let mid = (K + 1) / 2

var tree = Array(repeating: 0, count: 65539)
var arr = Array(repeating: 1, count: N + 1)
for i in 1...N {
    let num = Int(readLine()!)!;
    arr[i] = num + 1                                    // 0도 들어올 수 있는데, 펜윅트리에서 index는 양수여야 하므로 +1을 해서 기록해줌
}

for i in 0...K-1 {                                      // 일반화를 위해 맨 앞에 0이 추가로 있다고 가정
    add(at: arr[i], val: 1)
}

var result: Int64 = 0
for i in K...N {
    add(at: arr[i - K], val: -1)                    // arr[i - K]는 이번 부분 수열 범위에서 제외됨
    add(at: arr[i], val: 1)                         // arr[i]는 이번 부분 수열 범위에 추가됨
    result += Int64(lb(targetSum: mid) - 1)         // 작은 수 부터 카운트했을 때 합이 (K+1)/2개 이상이 되는 수가 중앙값
}
print(result)