Swift Advanced Data Structures: Segment Trees

Mastering Segment Trees with Swift

Segment trees are versatile data structures used for storing intervals of data and efficiently querying and updating them. They are particularly useful in scenarios where you need to perform range queries or updates on large datasets.

Example in Swift:


struct SegmentTree {
    var tree: [Int]
    var size: Int
    
    init(_ nums: [Int]) {
        size = nums.count
        tree = Array(repeating: 0, count: size * 4)
        buildTree(nums, 0, size - 1, 0)
    }
    
    mutating func buildTree(_ nums: [Int], _ start: Int, _ end: Int, _ index: Int) {
        if start == end {
            tree[index] = nums[start]
            return
        }
        let mid = (start + end) / 2
        buildTree(nums, start, mid, 2 * index + 1)
        buildTree(nums, mid + 1, end, 2 * index + 2)
        tree[index] = tree[2 * index + 1] + tree[2 * index + 2]
    }
    
    mutating func update(_ index: Int, _ val: Int, _ start: Int, _ end: Int, _ nodeIndex: Int) {
        if start == end {
            tree[nodeIndex] = val
            return
        }
        let mid = (start + end) / 2
        if index <= mid {
            update(index, val, start, mid, 2 * nodeIndex + 1)
        } else {
            update(index, val, mid + 1, end, 2 * nodeIndex + 2)
        }
        tree[nodeIndex] = tree[2 * nodeIndex + 1] + tree[2 * nodeIndex + 2]
    }
    
    func query(_ left: Int, _ right: Int, _ start: Int, _ end: Int, _ nodeIndex: Int) -> Int {
        if left <= start && right >= end {
            return tree[nodeIndex]
        }
        if right < start || left > end {
            return 0
        }
        let mid = (start + end) / 2
        return query(left, right, start, mid, 2 * nodeIndex + 1) + query(left, right, mid + 1, end, 2 * nodeIndex + 2)
    }
}

Explanation: In this example, we define a SegmentTree struct that can be used to build a segment tree from an array of integers. The buildTree function constructs the segment tree recursively by splitting the array into halves and summing up the values at each node. The update function allows for updating a specific element in the original array, while the query function enables range queries by traversing the segment tree and returning the sum of values within the specified range.

Conclusion: Segment trees are powerful data structures for efficiently handling interval-based queries and updates. By mastering their implementation in Swift, you can enhance your problem-solving skills and tackle a wide range of computational problems effectively.

Stay tuned for the next blog post where we'll dive into another advanced data structure in Swift!

COFPROG