SegmentTree

[lefex]

class ViewController: UIViewController {

override func viewDidLoad() {
    super.viewDidLoad()
    
    let datas = [6, 7, 4000, 2, 8]

// let seg = Segment(datas: datas, defaultValue: 0, merger: { $0 + $1 }) let seg = Segment(datas: datas, defaultValue: 0) { (a, b) -> Int in if a > b { return a } return b } seg.debug() if let ret = seg.query(ql: 0, qr: datas.count-1) { print("========= query[0-(datas.count-1)] = (ret)") }

    if let ret = seg.query(ql: 1, qr: 3) {
        print("========= query[\(1)-\(3)] = \(ret)")
    }
    seg.set(index: 3, value: 666)
    if let ret = seg.query(ql: 3, qr: 3) {
        print("========= query[\(3)-\(3)] = \(ret)")
    }
}

}

class Segment<T> { var datas: [T]! var trees: [T]! var merger: (_ a: T, _ b: T) -> T

init(datas: [T], defaultValue: T, merger: @escaping (_ a: T, _ b: T) -> T) {
    self.merger = merger
    self.datas = datas
    trees = Array(repeating: defaultValue, count: datas.count * 4)
    bulidSegment(l: 0, r: datas.count-1, curIndex: 0)
}

private func bulidSegment(l: Int, r: Int, curIndex: Int) {
    if l == r {
        trees[curIndex] = datas[l]
        return
    }
    
    let lchildIndex = leftChildIndex(index: curIndex)
    let rchildIndex = rightChildIndex(index: curIndex)
    let mid = l + (r - l) / 2
    bulidSegment(l: l, r: mid, curIndex: lchildIndex)
    bulidSegment(l: mid+1, r: r, curIndex: rchildIndex)
    
    trees[curIndex] = merger(trees[lchildIndex], trees[rchildIndex])
}

func query(ql: Int, qr: Int) -> T? {
    if ql < 0 || ql > qr || ql > datas.count - 1 || qr < 0 || qr > datas.count - 1 {
        return nil
    }
    return query(tIndex: 0, l: 0, r: datas.count - 1, ql: ql, qr: qr)
}

// [0-5] [0-2]
private func query(tIndex: Int, l: Int, r: Int, ql: Int, qr: Int) -> T? {
    print("tIndex=\(tIndex),l=\(l),r=\(r), ql=\(ql), qr=\(qr)")
    if l == ql && r == qr {
        return trees[tIndex]
    }
    
    let mid = l + (r - l) / 2
    let lchildIndex = leftChildIndex(index: tIndex)
    let rchildIndex = rightChildIndex(index: tIndex)
    
    if qr <= mid {
        // 都在 left
        let ret = query(tIndex: lchildIndex, l: ql, r: mid, ql: ql, qr: qr)
        return ret
    } else if ql >= mid + 1 {
        // 都在 right
        let ret = query(tIndex: rchildIndex, l: mid+1, r: r, ql: ql, qr: qr)
        return ret
    }
    
    // 既在 left 又在 right，分开查 left 和 right，查到结构后 merger
    let lresult = query(tIndex: lchildIndex, l: l, r: mid, ql: ql, qr: mid)
    let rresult = query(tIndex: rchildIndex, l: mid + 1, r: r, ql: mid + 1, qr: qr)
    if let lresult = lresult {
        if let rresult = rresult {
            return merger(lresult, rresult)
        }
    }
    return nil
}

func set(index: Int, value: T) {
    if index < 0 || index > datas.count {
        return
    }
    datas[index] = value
    set(index: index, value: value, tIndex: 0, l: 0, r: datas.count-1)
}

private func set(index: Int, value: T, tIndex: Int, l: Int, r: Int) {
    if l == r {
        trees[tIndex] = value
        return
    }
    let mid = l + (r-l) / 2
    let lchildIndex = leftChildIndex(index: tIndex)
    let rchildIndex = rightChildIndex(index: tIndex)
    if index <= mid {
        set(index: index, value: value, tIndex: lchildIndex, l: l, r: mid)
    } else if index >= mid+1 {
        set(index: index, value: value, tIndex: rchildIndex, l: mid + 1, r: r)
    }
    trees[tIndex] = merger(trees[lchildIndex], trees[rchildIndex])
}

func leftChildIndex(index: Int) -> Int {
    return 2 * index + 1;
}

func rightChildIndex(index: Int) -> Int {
    return 2 * index + 2;
}

func debug() {
    for item in trees {
        print(item)
    }
}

}