백준 18436번 수열과 쿼리 37 - 스위프트(Swift) 풀이

 수열과 쿼리 시리즈 중에 하나. 변화하는 수열에서 구간 내 짝수/홀수 개수를 묻는 쿼리를 여러 번 수행해야 한다. 쿼리 종류가 2가지이지만, 결국 한 가지인 거랑 똑같은데, 짝수 아니면 홀수니까 둘 중 하나만 카운트하는 세그먼트 트리를 두고, 만약 반대를 물어보면 구간 길이에서 물어본 것의 반대 개수만큼을 빼주면 된다.

 짝수/홀수 중 한 종류만 카운트하게 두었기 때문에 update 작업을 할 때 조심해야 한다. 짝수 -> 짝수 또는 홀수 -> 홀수로의 업데이트를 할 때는 결국 변화가 없는 것과 같다. (구체적인 값은 중요하지 않고 짝수/홀수 여부만 상관있기 때문에.) 따라서 만약 짝수를 카운트하는 세그먼트 트리를 뒀다면, 짝수 -> 홀수로 변하면 -1을 해주고, 홀수 -> 짝수로 변하면 +1을 해주면 된다. 홀수를 카운트하도록 구현한 경우에는 반대로 하면 되겠다.

 나는 짝수를 카운트하는 펜윅트리를 두는 방식으로 풀었다.

import Foundation

class FenwickTree {
    var tree = Array(repeating: 0, count: 100001)
    
    init(arr: [Int]) {
        for (index, num) in arr.enumerated() {
            if num % 2 == 0 {
                add(at: index + 1, 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)
        }
    }
}

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

var fwt = FenwickTree(arr: arr)     // 짝수만 카운트
var result = ""

let M = Int(readLine()!)!
for _ in 0..<M {
    let input = readLine()!.split(separator: " ").map{ Int(String($0))! }
    
    if input[0] == 1 {
        let index = input[1]
        let newVal = input[2]
        
        if arr[index - 1] % 2 != 0 && newVal % 2 == 0 {         // 홀수 > 짝수로 업데이트 : 짝수 개수 + 1
            fwt.add(at: index, val: 1)
        } else if arr[index - 1] % 2 == 0 && newVal % 2 != 0 {  // 짝수 > 홀수로 업데이트 : 짝수 개수 - 1
            fwt.add(at: index, val: -1)
        }
        
        arr[index - 1] = newVal
    } else {
        let lo = input[1]
        let hi = input[2]
        
        let evenCount = fwt.sum(to: hi) - fwt.sum(to: lo - 1)   // 짝수 개수 = 구간합 [lo...hi]
        let oddCount = hi - lo + 1 - evenCount                  // 홀수 개수 = 구간 길이 - 짝수 개수
        
        result.write("\(input[0] == 2 ? evenCount : oddCount)\n")
    }
}

print(result)