Swift Coding Test Course, Segment Tree

Hello! In this lecture, we will take a closer look at the segment tree, which is one of the data structures. The segment tree is a powerful data structure that can efficiently handle interval queries, especially useful for calculating the range sum, range minimum, and range maximum of an array. In this text, we will explore the basic concepts of segment trees, how to implement them, and solve problems that are frequently asked in practice.

1. What is a Segment Tree?

A segment tree is a powerful tool for efficiently managing the interval information of an array. Typically, when the size of the array is N, a segment tree uses approximately 2 * 2⌈log₂(N)⌉ memory space. This is due to the use of a complete binary tree structure. Fundamentally, the segment tree can perform the following two operations efficiently:

  • Interval Query: Quickly retrieve information about a specific interval.
  • Update: Rapidly update the interval information affected after changing a specific element of the array.

2. Structure of the Segment Tree

The segment tree is made up of nodes, each representing a specific interval. For example, there is a root node that manages the interval from index 0 to N-1 of the array, and this node divides the array into two halves through two child nodes. By continuously dividing the array in this manner, each node holds the information of a specific interval.

2.1 Node Definition

Each node contains the following information:

  • Start Index (start): The starting point of the interval
  • End Index (end): The ending point of the interval
  • Value (value): A variable to store the information of the interval (e.g., sum, minimum value, etc.)

3. Creating and Querying the Segment Tree

To implement a segment tree, we first need to build the tree. For this, we use a method that recursively creates segment tree nodes based on the input array. As a simple example, let’s create a segment tree for range sums.


// Swift Code
class SegmentTree {
    var tree: [Int] // Array to store the segment tree
    var n: Int // Size of the array

    init(_ data: [Int]) {
        self.n = data.count
        self.tree = Array(repeating: 0, count: 4 * n) // Initialize the tree array
        buildTree(data: data, node: 1, start: 0, end: n - 1)
    }

    // Build the segment tree using the interval of the array
    func buildTree(data: [Int], node: Int, start: Int, end: Int) {
        if start == end {
            tree[node] = data[start] // Store value in leaf node
        } else {
            let mid = (start + end) / 2
            buildTree(data: data, node: 2 * node, start: start, end: mid) // Left subtree
            buildTree(data: data, node: 2 * node + 1, start: mid + 1, end: end) // Right subtree
            tree[node] = tree[2 * node] + tree[2 * node + 1] // Update parent node value
        }
    }
}

4. Processing Segment Tree Queries

The segment tree now needs to add a feature to calculate range sums. To handle range sum queries, we will add the following function:


// Function to calculate the sum of a given interval
func query(node: Int, start: Int, end: Int, l: Int, r: Int) -> Int {
    if r < start || end < l { // If intervals do not overlap
        return 0 // Default value
    }
    if l <= start && end <= r { // If the interval is completely contained
        return tree[node]
    }
    let mid = (start + end) / 2
    let leftSum = query(node: 2 * node, start: start, end: mid, l: l, r: r) // Query left subtree
    let rightSum = query(node: 2 * node + 1, start: mid + 1, end: end, l: l, r: r) // Query right subtree
    return leftSum + rightSum // Sum the results
}

5. Updating the Segment Tree

We will also add a function to update the elements of the array. When changing a specific index of the array, here’s how we can quickly update the information in the segment tree:


// Function to update a specific index of the array
func update(node: Int, start: Int, end: Int, idx: Int, value: Int) {
    if start == end { // Reached leaf node
        tree[node] = value // Update the node
    } else {
        let mid = (start + end) / 2
        if start <= idx && idx <= mid {
            update(node: 2 * node, start: start, end: mid, idx: idx, value: value) // Update left subtree
        } else {
            update(node: 2 * node + 1, start: mid + 1, end: end, idx: idx, value: value) // Update right subtree
        }
        tree[node] = tree[2 * node] + tree[2 * node + 1] // Update parent node value
    }
}

6. Example of a Practical Problem

Now that we have looked at the basic structure of the segment tree and how to perform queries and updates, let's tackle a problem that is frequently asked in practice. The problem is as follows:

Problem Description

Given an array, answer the following questions:

  1. Calculate the sum of the interval [L, R].
  2. Update the value of index I to V.

Input Format

N (array size)
arr[0], arr[1], ..., arr[N-1]
Q (number of queries)
Each query is given in the following format:
1 L R (range sum query)
2 I V (update query)

Output Format

Output for each range sum query

Example

5
1 2 3 4 5
3
1 1 3
2 2 10
1 1 3
Example Output
6

7. Problem-Solving Process

  1. Build a segment tree for the input array.
  2. Depending on the query type, call the query() function for range queries and the update() function for update queries.
  3. Print the results.

// Main function to solve the problem
import Foundation

func main() {
    // Input processing
    let n = Int(readLine()!)!
    let arr = readLine()!.split(separator: " ").map { Int($0)! }
    let q = Int(readLine()!)!

    // Build the segment tree
    let segmentTree = SegmentTree(arr)

    // Process queries
    for _ in 0..

8. Conclusion

Through this lecture on segment trees, we learned the importance of data structures and how to utilize them. Segment trees can be effectively used for various interval query problems. We hope you will enhance your skills in applying segment trees through more practice problems. Thank you!