Skip to content

Commit d5cba78

Browse files
hollancehollance
committed
Merge pull request kodecocodes#34 from goingreen/master
Add segment tree
2 parents 1950f48 + 00e5bc6 commit d5cba78

File tree

6 files changed

+308
-0
lines changed

6 files changed

+308
-0
lines changed

‎Segment Tree/Images/Structure.png

76.4 KB
Loading

‎Segment Tree/README.markdown

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
# Segment Tree
2+
3+
I'm pleased to present to you Segment Tree. It's actually one of my favorite data structures because it's very flexible and simple in realization.
4+
5+
Let's suppose that you have array **a** of some type and some associative function _**f**_(e.g. sum, multiplication, min, max, gcd).
6+
Your task is:
7+
* answer a queries for given **l** and **r**: `f(a[l], a[l+1], ..., a[r-1], a[r])`
8+
* support replacing item at some index `a[index] = newItem`
9+
10+
Here's naive approach if our array's type is Int and _**f**_ is just sum of two integers:
11+
```swift
12+
func query(array: [Int], l: Int, r: Int) -> Int {
13+
var sum = 0
14+
for i in l...r {
15+
sum += array[i]
16+
}
17+
return sum
18+
}
19+
```
20+
The running time of this algorithm is **O(n)** in worst case (**l = 0, r = n-1**). And if we have **m** queries to answer we get **O(m*n)** complexity.
21+
If we have **n = 10^5** and **m = 100** our algorithm will do **10^7** units of work, ouh it's sounds not very good. Let's look how we can improve it.
22+
23+
Segment trees allow us answer a queries and replace items on **O(log n)**, isn't it magic?:sparkles:
24+
The main idea of segment trees is simple: we precalculate some segments in our array and then we can use it without repeating calculation.
25+
## Structure of segment tree
26+
27+
Segment tree is just [binary tree](../Binary Tree/) where each node has:
28+
* `leftBound`
29+
* `rightBound`
30+
* `value` is actually `f(a[leftBound], a[leftBound+1], .., a[rightBound])`
31+
* `leftChild`
32+
* `rightChild`
33+
34+
Here's structure of segment tree for given array `[1, 2, 3, 4]` and **f = a+b**, pairs of leftBound and rightBound marked in red
35+
36+
![structure](Images/Structure.png)
37+
38+
## Building segment tree
39+
40+
Let's see how to build node of segment tree.
41+
```swift
42+
init(array: [T], leftBound: Int, rightBound: Int, function: (T, T) -> T) {
43+
self.leftBound = leftBound
44+
self.rightBound = rightBound
45+
self.function = function
46+
if leftBound == rightBound {
47+
value = array[leftBound]
48+
} else {
49+
let middle = (leftBound + rightBound) / 2
50+
leftChild = SegmentTree<T>(array: array, leftBound: leftBound, rightBound: middle, function: function)
51+
rightChild = SegmentTree<T>(array: array, leftBound: middle+1, rightBound: rightBound, function: function)
52+
value = function(leftChild!.value, rightChild!.value)
53+
}
54+
}
55+
```
56+
57+
If out current leftBound and rightBound are equal it means that we are in leaf so we don't have any child nodes and we just fill in `value` property with `array[leftBound]` else we have two child nodes. In that case we divide our current segment into two equal (if length is even) segments: `middle = (leftBound + rightBound) / 2` **[leftBound, middle]** and **[middle+1, rightBound]** and then we build our child nodes for that segments. After we build our child nodes we can easily calculate our value as `value = function(leftChild!.value, rightChild!.value)` because **f(leftBound, rightBound) = f(f(leftBound, middle), f(middle+1, rightBound))**
58+
59+
## Getting answer to query
60+
61+
```swift
62+
public func queryWithLeftBound(leftBound: Int, rightBound: Int) -> T {
63+
if self.leftBound == self.rightBound {
64+
return self.value
65+
} else if leftChild!.rightBound < leftBound {
66+
return rightChild!.queryWithLeftBound(leftBound, rightBound: rightBound)
67+
} else if rightChild!.leftBound > rightBound {
68+
return leftChild!.queryWithLeftBound(leftBound, rightBound: rightBound)
69+
} else {
70+
let leftResult = leftChild!.queryWithLeftBound(leftBound, rightBound: leftChild!.rightBound)
71+
let rightResult = rightChild!.queryWithLeftBound(rightChild!.leftBound, rightBound: rightBound)
72+
return function(leftResult, rightResult)
73+
}
74+
}
75+
```
76+
Firstly, we check if our current node is leaf, if it is we just return its value else
77+
we check that our query segment fully lies in rightChild, if so we return result of query on rightChild else if segment lies in leftChild we return result of query on leftChild. If our query lies in both child then we combine results of query on both child.
78+
79+
## Replacing items
80+
81+
```swift
82+
public func replaceItemAtIndex(index: Int, withItem item: T) {
83+
if leftBound == rightBound {
84+
value = item
85+
} else {
86+
if leftChild!.rightBound >= index {
87+
leftChild!.replaceItemAtIndex(index, withItem: item)
88+
} else {
89+
rightChild!.replaceItemAtIndex(index, withItem: item)
90+
}
91+
value = function(leftChild!.value, rightChild!.value)
92+
}
93+
}
94+
```
95+
As always we check if current node is leaf and if so we just change its value otherwise we need to find out to which child index belongs. After that we call same function on that child and after that we recalculate our value because it needs to be in actual state.
96+
97+
## Examples
98+
99+
Examples of using segment trees can be found in playground
100+
101+
*Written for Swift Algorithm Club by [Artur Antonov](https://github.com/goingreen)*

‎Segment Tree/SegmentTree.playground/Contents.swift

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
//: Playground - noun: a place where people can play
2+
3+
public class SegmentTree<T> {
4+
5+
private var value: T
6+
private var function: (T, T) -> T
7+
private var leftBound: Int
8+
private var rightBound: Int
9+
private var leftChild: SegmentTree<T>?
10+
private var rightChild: SegmentTree<T>?
11+
12+
13+
init(array: [T], leftBound: Int, rightBound: Int, function: (T, T) -> T) {
14+
self.leftBound = leftBound
15+
self.rightBound = rightBound
16+
self.function = function
17+
if leftBound == rightBound {
18+
value = array[leftBound]
19+
} else {
20+
let middle = (leftBound + rightBound) / 2
21+
leftChild = SegmentTree<T>(array: array, leftBound: leftBound, rightBound: middle, function: function)
22+
rightChild = SegmentTree<T>(array: array, leftBound: middle+1, rightBound: rightBound, function: function)
23+
value = function(leftChild!.value, rightChild!.value)
24+
}
25+
}
26+
27+
convenience init(array: [T], function: (T, T) -> T) {
28+
self.init(array: array, leftBound: 0, rightBound: array.count-1, function: function)
29+
}
30+
31+
public func queryWithLeftBound(leftBound: Int, rightBound: Int) -> T {
32+
if self.leftBound == self.rightBound {
33+
return self.value
34+
} else if leftChild!.rightBound < leftBound {
35+
return rightChild!.queryWithLeftBound(leftBound, rightBound: rightBound)
36+
} else if rightChild!.leftBound > rightBound {
37+
return leftChild!.queryWithLeftBound(leftBound, rightBound: rightBound)
38+
} else {
39+
let leftResult = leftChild!.queryWithLeftBound(leftBound, rightBound: leftChild!.rightBound)
40+
let rightResult = rightChild!.queryWithLeftBound(rightChild!.leftBound, rightBound: rightBound)
41+
return function(leftResult, rightResult)
42+
}
43+
}
44+
45+
public func replaceItemAtIndex(index: Int, withItem item: T) {
46+
if leftBound == rightBound {
47+
value = item
48+
} else {
49+
if leftChild!.rightBound >= index {
50+
leftChild!.replaceItemAtIndex(index, withItem: item)
51+
} else {
52+
rightChild!.replaceItemAtIndex(index, withItem: item)
53+
}
54+
value = function(leftChild!.value, rightChild!.value)
55+
}
56+
}
57+
}
58+
59+
60+
61+
62+
let array = [1, 2, 3]
63+
64+
let sumSegmentTree = SegmentTree(array: array, function: +)
65+
66+
print(sumSegmentTree.queryWithLeftBound(0, rightBound: 2)) // 1 + 2 + 3 + 4 = 10
67+
print(sumSegmentTree.queryWithLeftBound(1, rightBound: 2)) // 2 + 3 = 5
68+
print(sumSegmentTree.queryWithLeftBound(0, rightBound: 0)) // 1 = 1
69+
70+
sumSegmentTree.replaceItemAtIndex(0, withItem: 2) //our array now is [2, 2, 3, 4]
71+
72+
print(sumSegmentTree.queryWithLeftBound(0, rightBound: 0)) // 2 = 2
73+
print(sumSegmentTree.queryWithLeftBound(0, rightBound: 1)) // 2 + 2 = 4
74+
75+
76+
//you can use any associative function (i.e (a+b)+c == a+(b+c)) as function for segment tree
77+
func gcd(m: Int, _ n: Int) -> Int {
78+
var a = 0
79+
var b = max(m, n)
80+
var r = min(m, n)
81+
82+
while r != 0 {
83+
a = b
84+
b = r
85+
r = a % b
86+
}
87+
return b
88+
}
89+
90+
let gcdArray = [2, 4, 6, 3, 5]
91+
92+
let gcdSegmentTree = SegmentTree(array: gcdArray, function: gcd)
93+
94+
print(gcdSegmentTree.queryWithLeftBound(0, rightBound: 1)) // gcd(2, 4) = 2
95+
print(gcdSegmentTree.queryWithLeftBound(2, rightBound: 3)) // gcd(6, 3) = 3
96+
print(gcdSegmentTree.queryWithLeftBound(1, rightBound: 3)) // gcd(4, 6, 3) = 1
97+
print(gcdSegmentTree.queryWithLeftBound(0, rightBound: 4)) // gcd(2, 4, 6, 3, 5) = 1
98+
99+
gcdSegmentTree.replaceItemAtIndex(3, withItem: 10) //gcdArray now is [2, 4, 6, 10, 5]
100+
101+
print(gcdSegmentTree.queryWithLeftBound(3, rightBound: 4)) // gcd(10, 5) = 5
102+
103+
104+
//example of segment tree which finds minimum on given range
105+
let minArray = [2, 4, 1, 5, 3]
106+
107+
let minSegmentTree = SegmentTree(array: minArray, function: min)
108+
109+
print(minSegmentTree.queryWithLeftBound(0, rightBound: 4)) // min(2, 4, 1, 5, 3) = 1
110+
print(minSegmentTree.queryWithLeftBound(0, rightBound: 1)) // min(2, 4) = 2
111+
112+
minSegmentTree.replaceItemAtIndex(2, withItem: 10) // minArray now is [2, 4, 10, 5, 3]
113+
114+
print(minSegmentTree.queryWithLeftBound(0, rightBound: 4)) // min(2, 4, 10, 5, 3) = 2
115+
116+
117+
//type of elements in array can be any type which has some associative function
118+
let stringArray = ["a", "b", "c", "A", "B", "C"]
119+
120+
let stringSegmentTree = SegmentTree(array: stringArray, function: +)
121+
122+
print(stringSegmentTree.queryWithLeftBound(0, rightBound: 1)) // "a"+"b" = "ab"
123+
print(stringSegmentTree.queryWithLeftBound(2, rightBound: 3)) // "c"+"A" = "cA"
124+
print(stringSegmentTree.queryWithLeftBound(1, rightBound: 3)) // "b"+"c"+"A" = "bcA"
125+
print(stringSegmentTree.queryWithLeftBound(0, rightBound: 5)) // "a"+"b"+"c"+"A"+"B"+"C" = "abcABC"
126+
127+
stringSegmentTree.replaceItemAtIndex(0, withItem: "I")
128+
stringSegmentTree.replaceItemAtIndex(1, withItem: " like")
129+
stringSegmentTree.replaceItemAtIndex(2, withItem: " algorithms")
130+
stringSegmentTree.replaceItemAtIndex(3, withItem: " and")
131+
stringSegmentTree.replaceItemAtIndex(4, withItem: " swift")
132+
stringSegmentTree.replaceItemAtIndex(5, withItem: "!")
133+
134+
print(stringSegmentTree.queryWithLeftBound(0, rightBound: 5))
135+

‎Segment Tree/SegmentTree.playground/contents.xcplayground

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
2+
<playground version='5.0' target-platform='osx'>
3+
<timeline fileName='timeline.xctimeline'/>
4+
</playground>

‎Segment Tree/SegmentTree.playground/timeline.xctimeline

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<Timeline
3+
version = "3.0">
4+
<TimelineItems>
5+
</TimelineItems>
6+
</Timeline>

‎Segment Tree/SegmentTree.swift

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
/*
2+
Segment tree
3+
build is O(n)
4+
query is O(log n)
5+
replace item is O(log n)
6+
*/
7+
8+
public class SegmentTree<T> {
9+
10+
private var value: T
11+
private var function: (T, T) -> T
12+
private var leftBound: Int
13+
private var rightBound: Int
14+
private var leftChild: SegmentTree<T>?
15+
private var rightChild: SegmentTree<T>?
16+
17+
18+
init(array: [T], leftBound: Int, rightBound: Int, function: (T, T) -> T) {
19+
self.leftBound = leftBound
20+
self.rightBound = rightBound
21+
self.function = function
22+
if leftBound == rightBound {
23+
value = array[leftBound]
24+
} else {
25+
let middle = (leftBound + rightBound) / 2
26+
leftChild = SegmentTree<T>(array: array, leftBound: leftBound, rightBound: middle, function: function)
27+
rightChild = SegmentTree<T>(array: array, leftBound: middle+1, rightBound: rightBound, function: function)
28+
value = function(leftChild!.value, rightChild!.value)
29+
}
30+
}
31+
32+
convenience init(array: [T], function: (T, T) -> T) {
33+
self.init(array: array, leftBound: 0, rightBound: array.count-1, function: function)
34+
}
35+
36+
public func queryWithLeftBound(leftBound: Int, rightBound: Int) -> T {
37+
if self.leftBound == self.rightBound {
38+
return self.value
39+
} else if leftChild!.rightBound < leftBound {
40+
return rightChild!.queryWithLeftBound(leftBound, rightBound: rightBound)
41+
} else if rightChild!.leftBound > rightBound {
42+
return leftChild!.queryWithLeftBound(leftBound, rightBound: rightBound)
43+
} else {
44+
let leftResult = leftChild!.queryWithLeftBound(leftBound, rightBound: leftChild!.rightBound)
45+
let rightResult = rightChild!.queryWithLeftBound(rightChild!.leftBound, rightBound: rightBound)
46+
return function(leftResult, rightResult)
47+
}
48+
}
49+
50+
public func replaceItemAtIndex(index: Int, withItem item: T) {
51+
if leftBound == rightBound {
52+
value = item
53+
} else {
54+
if leftChild!.rightBound >= index {
55+
leftChild!.replaceItemAtIndex(index, withItem: item)
56+
} else {
57+
rightChild!.replaceItemAtIndex(index, withItem: item)
58+
}
59+
value = function(leftChild!.value, rightChild!.value)
60+
}
61+
}
62+
}

0 commit comments

Comments
 (0)