lib: move some examples/* into specific subdirectories of their lib
[nit.git] / lib / trees / bintree.nit
index 8165a24..b689700 100644 (file)
@@ -45,6 +45,8 @@ class BinTreeMap[K: Comparable, E]
        redef type N: BinTreeNode[K, E]
 
        private var len = 0
+       private var first_node: nullable BinTreeNode[K, E] = null
+       private var last_node: nullable BinTreeNode[K, E] = null
 
        # O(n) in worst case, average is O(h) with h: tree height
        #
@@ -53,22 +55,48 @@ class BinTreeMap[K: Comparable, E]
        #     tree[1] = "n1"
        #     assert not tree.is_empty
        redef fun is_empty do return root == null
+
+       # O(n) in worst case, average is O(h) with h: tree height
+       #
+       #     var tree = new BinTreeMap[Int, String]
+       #     assert not tree.has_key(1)
+       #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
+       #     assert not tree.has_key(0)
+       #     assert tree.has_key(2)
+       #     assert not tree.has_key(6)
+       redef fun has_key(key: K): Bool do
+               if is_empty then return false
+               var res = search_down(root.as(not null), key)
+               if res != null then
+                       cache_node = res
+                       return true
+               end
+               return false
+       end
+
+       private var cache_node: nullable N = null
+
        # Get the node value associated to `key`
        # O(n) in worst case, average is O(h) with h: tree height
        #
        #     var tree = new BinTreeMap[Int, String]
        #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
+       #     assert tree.has_key(1)
        #     assert tree[1] == "n1"
+       #     assert tree.has_key(1)
+       #     assert tree[2] == "n2"
        redef fun [](key: K): E do
-               assert not_empty: root != null
+               assert not_empty: not is_empty
+               if cache_node != null and cache_node.key == key then return cache_node.value
                var res = search_down(root.as(not null), key)
                assert has_key: res != null
                return res.value
        end
 
        protected fun search_down(from: N, key: K): nullable N do
-               if key == from.key then return from
-               if from.left != null and key < from.key then
+               var cmp = key <=> from.key
+               if cmp == 0 then return from
+               if from.left != null and cmp < 0 then
                        return search_down(from.left.as(not null), key)
                else if from.right != null then
                        return search_down(from.right.as(not null), key)
@@ -127,18 +155,27 @@ class BinTreeMap[K: Comparable, E]
                else
                        shift_down(root.as(not null), node)
                end
+               if first_node == null then
+                       first_node = node
+               end
+               if last_node != null then
+                       last_node.next = node
+                       node.prev = last_node
+               end
+               last_node = node
        end
 
        # Push down the `node` in tree from a specified `from` index
        protected fun shift_down(from, node: N) do
-               if node.key < from.key then
+               var cmp = node.key <=> from.key
+               if cmp < 0 then
                        if from.left == null then
                                from.left = node
                                node.parent = from
                        else
                                shift_down(from.left.as(not null), node)
                        end
-               else if node.key > from.key then
+               else if cmp > 0 then
                        if from.right == null then
                                from.right = node
                                node.parent = from
@@ -178,6 +215,16 @@ class BinTreeMap[K: Comparable, E]
                        min.left = node.left
                        min.left.parent = min
                end
+               if first_node == node then
+                       first_node = null
+               end
+               if last_node == node then
+                       last_node = node.prev
+                       last_node.next = null
+               else
+                       node.prev.next = node.next
+                       node.next.prev = node.prev
+               end
                return node.value
        end
 
@@ -302,12 +349,27 @@ class BinTreeMap[K: Comparable, E]
        #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
        #     assert tree.length == 5
        redef fun length do return len
+
+       # Nodes are iterated in the same order in which they were added to the tree.
+       # O(n)
+       #
+       #     var tree = new BinTreeMap[Int, String]
+       #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
+       #     var keys = new Array[Int]
+       #     for k, v in tree do
+       #         keys.add k
+       #     end
+       #     assert keys == [4, 2, 1, 5, 3]
+       redef fun iterator do return new BinTreeMapIterator[K, E](self)
 end
 
 # TreeNode used by BinTree
 class BinTreeNode[K: Comparable, E]
        super TreeNode[K, E]
 
+       private var prev: nullable BinTreeNode[K, E]
+       private var next: nullable BinTreeNode[K, E]
+
        redef type SELF: BinTreeNode[K, E]
 
        init(key: K, item: E) do
@@ -322,7 +384,7 @@ class BinTreeNode[K: Comparable, E]
        # set `left` child for this node (or null if left no child)
        # ENSURE: node.key < key (only if node != null)
        fun left=(node: nullable SELF) do
-               assert node != null implies node.key < key
+               #assert node != null implies node.key < key
                left_node = node
        end
 
@@ -334,9 +396,7 @@ class BinTreeNode[K: Comparable, E]
        # set `right` child for this node (or null if right no child)
        # ENSURE: node.key < key (only if node != null)
        fun right=(node: nullable SELF) do
-               if node != null then
-                       assert node.key > key
-               end
+               #assert node != null implies node.key > key
                right_node = node
        end
 
@@ -378,6 +438,20 @@ class BinTreeNode[K: Comparable, E]
                end
        end
 
-       redef fun to_s do return "\{{key}: {value}\}"
+       redef fun to_s do return "\{{key}: {value or else ""}\}"
 end
 
+private class BinTreeMapIterator[K: Comparable, E]
+       super MapIterator[K, E]
+
+       var current: nullable BinTreeNode[K, E]
+
+       init(tree: BinTreeMap[K, E]) do
+               current = tree.first_node
+       end
+
+       redef fun is_ok do return not current == null
+       redef fun next do current = current.next
+       redef fun item do return current.value
+       redef fun key do do return current.key
+end