-
Notifications
You must be signed in to change notification settings - Fork 1
/
GeoTrie.scala
98 lines (89 loc) · 3.27 KB
/
GeoTrie.scala
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
package damysos
import scala.annotation.tailrec
import cats.implicits._
import Constants._
protected sealed trait GeoTrie
protected final case class Leaf(locations: Array[PointOfInterest] = Array.empty) extends GeoTrie
protected final case class Node(
private val children: Array[Array[GeoTrie]] = Array.fill(TreeBreadth)(Array.ofDim(TreeBreadth)),
) extends GeoTrie {
def toArray: Array[PointOfInterest] =
children.foldLeft(Array.empty[PointOfInterest])((acc, arr) =>
arr.foldLeft(acc)((acc2, geoTrie) =>
geoTrie match {
case node: Node => acc2 concat node.toArray
case Leaf(locations) => acc2 concat locations
}
)
)
def size: Int =
children.foldLeft(0)((acc, arr) =>
arr.foldLeft(acc)((acc2, geoTrie) =>
geoTrie match {
case node: Node => acc2 + node.size
case Leaf(location) => acc2 + location.size
}
)
)
@tailrec
def findLeaf(path: Array[(Int, Int)], trie: GeoTrie = this): Option[GeoTrie] =
if (path.length == 1)
trie match {
case Node(children) => {
val (latIndex, longIndex) = path.head
if (children(latIndex).isDefinedAt(longIndex))
Some(children(latIndex)(longIndex))
else None
}
case _ => None
}
else
trie match {
case Node(children) => {
val (latIndex, longIndex) = path.head
if (children(latIndex).isDefinedAt(longIndex))
findLeaf(path.tail, children(latIndex)(longIndex))
else None
}
case _ => None
}
// Helper function to update a node's child at the given coordinates.
private def nodeItemUpdate(node: Node, latIndex: Int, longIndex: Int, item: GeoTrie): Node =
node.copy(
children=node.children.updated(
latIndex,
node.children(latIndex).updated(longIndex, item),
)
)
// Descends the given path (creates it if necessary), and applies the given function to the Leaf
// at the end of that path (or the Leaf it created there, if the path didn't exist fully).
private def updateAtPath(fn: Array[PointOfInterest] => Array[PointOfInterest])
(path: Array[(Int, Int)], node: Node): Node =
if (path.length === 1) {
val (latIndex, longIndex) = path.head
val leaf = {
if (children(latIndex).isDefinedAt(longIndex))
node.children(latIndex)(longIndex) match {
case leaf: Leaf => leaf
case _ => Leaf()
}
else Leaf()
}
nodeItemUpdate(node, latIndex, longIndex, leaf.copy(fn(leaf.locations)))
} else {
val (latIndex, longIndex) = path.head
val subNode = {
if (node.children(latIndex).isDefinedAt(longIndex))
node.children(latIndex)(longIndex) match {
case node: Node => node
case _ => Node()
}
else Node()
}
nodeItemUpdate(node, latIndex, longIndex, updateAtPath(fn)(path.tail, subNode))
}
def insertAtPath(item: PointOfInterest, path: Array[(Int, Int)], node: Node = this): Node =
updateAtPath(_ :+ item)(path, node)
def removeAtPath(item: PointOfInterest, path: Array[(Int, Int)], node: Node = this): Node =
updateAtPath(_.filter(_ != item))(path, node)
}