Red black tree sort/Wren

From Rosetta Code

This is based, more or less, on the code here which I thought was reasonably intelligible.

Although red black trees don't necessarily need to disallow duplicate keys, I've adjusted the code to optionally disallow them. Not that it matters here as the method of constructing the random keys ensures that there will be no duplicates in any case.

import "random" for Random

/* Represents a node in the red-black tree. */
class Node {
    // constructs a new node
    construct new(data, parent, left, right, color) {
        _data = data
        _parent = parent
        _left = left
        _right = right
        _color = color  // 0 black, 1 red
    }

    // constructs a new 'sentinel' node
    static new() { new(null, null, null, null, 0) }

    data   { _data   }
    parent { _parent }
    left   { _left   }
    right  { _right  }
    color  { _color  }

    data=(d)   { _data = d   }
    parent=(p) { _parent = p }
    left=(l)   { _left = l   }
    right=(r)  { _right = r  }
    color=(c)  { _color = c  }
}

/* Represents a red-black tree data structure. */
class RBTree {
    // constructs a new red-black tree specifying whether duplicate keys are allowed or not
    construct new(allowDups) {
        _tnull = Node.new()
        _root = _tnull
        _allowDups = allowDups
    }

    // convenience version of the constructor which always allows duplicate keys
    static new() { new(true) }

    // returns the root node
    root { _root }

    /* private helper methods for basic operations */

    preOrderHelper_(node) {
        if (node != _tnull) {
            System.print(node.data + " ")
            preOrderHelper_(node.left)
            preOrderHelper_(node.right)
        }
    }

    inOrderHelper_(node) {
        if (node != _tnull) {
            inOrderHelper_(node.left)
            System.print(node.data + " ")
            inOrderHelper_(node.right)
        }
    }

    postOrderHelper_(node) {
        if (node != _tnull) {
            postOrderHelper_(node.left)
            postOrderHelper_(node.right)
            System.print(node.data + " ")
        }
    }

    searchTreeHelper_(node, key) {
        if (node == _tnull || key == node.data) return node
        if (key < node.data) return searchTreeHelper_(node.left, key)
        return searchTreeHelper_(node.right, key)
    }

    fixDelete_(x) {
        var s
        while (x != _root && x.color == 0) {
            if (x == x.parent.left) {
                s = x.parent.right
                if (s.color == 1) {
                    s.color = 0
                    x.parent.color = 1
                    leftRotate(x.parent)
                    s = x.parent.right
                }
                if (s.left.color == 0 && s.right.color == 0) {
                    s.color = 1
                    x = x.parent
                } else {
                    if (s.right.color == 0) {
                        s.left.color = 0
                        s.color = 1
                        rightRotate(s)
                        s = x.parent.right
                    }
                    s.color = x.parent.color
                    x.parent.color = 0
                    s.right.color = 0
                    leftRotate(x.parent)
                    x = _root
                }
            } else {
                s = x.parent.left
                if (s.color == 1) {
                    s.color = 0
                    x.parent.color = 1
                    rightRotate(x.parent)
                    s = x.parent.left
                }
                if (s.left.color == 0 && s.right.color == 0) {
                    s.color = 1
                    x = x.parent
                } else {
                    if (s.left.color == 0) {
                        s.right.color = 0
                        s.color = 1
                        leftRotate(s)
                        s = x.parent.left
                    }
                    s.color = x.parent.color
                    x.parent.color = 0
                    s.left.color = 0
                    rightRotate(x.parent)
                    x = _root
                }
            }
        }
        x.color = 0
    }

    rbTransplant_(u, v) {
        if (!u.parent) {
            _root = v
        } else if (u == u.parent.left) {
            u.parent.left = v
        } else {
            u.parent.right = v
        }
        v.parent = u.parent
    }

    deleteNodeHelper_(node, key) {
        // find the node containing key
        var z = _tnull
        var x
        var y
        while (node != _tnull) {
            if (node.data == key) z = node
            if (node.data <= key) {
                node = node.right
            } else {
                node = node.left
            }
        }
        if (z == _tnull) {
            System.print("Couldn't find key in the tree.")
            return
        }
        y = z
        var yOriginalColor = y.color
        if (z.left == _tnull) {
            x = z.right
            rbTransplant_(z, z.right)
        } else if (z.right == _tnull) {
            x = z.left
            rbTransplant_(z, z.left)
        } else {
            y = minimum(z.right)
            yOriginalColor = y.color
            x = y.right
            if (y.parent == z) {
                x.parent = y
            } else {
                rbTransplant_(y, y.right)
                y.right = z.right
                y.right.parent = y
            }
            rbTransplant_(z, y)
            y.left = z.left
            y.left.parent = y
            y.color = z.color
        }
        if (yOriginalColor == 0) fixDelete_(x)
    }

    fixInsert_(k) {
        var u
        while (k.parent.color == 1) {
            if (k.parent == k.parent.parent.right) {
                u = k.parent.parent.left // uncle
                if (u.color == 1) {
                    u.color = 0
                    k.parent.color = 0
                    k.parent.parent.color = 1
                    k = k.parent.parent
                } else {
                    if (k == k.parent.left) {
                        k = k.parent
                        rightRotate(k)
                    }
                    k.parent.color = 0
                    k.parent.parent.color = 1
                    leftRotate(k.parent.parent)
                }
            } else {
                u = k.parent.parent.right // uncle
                if (u.color == 1) {
                    u.color = 0
                    k.parent.color = 0
                    k.parent.parent.color = 1
                    k = k.parent.parent
                } else {
                    if (k == k.parent.right) {
                        k = k.parent
                        leftRotate(k)
                    }
                    k.parent.color = 0
                    k.parent.parent.color = 1
                    rightRotate(k.parent.parent)
                }
            }
            if (k == _root) break
        }
        _root.color = 0
    }

    printHelper_(root, indent, last) {
        if (root != _tnull) {
           System.write(indent)
           if (last) {
              System.write("R----")
              indent = indent + "     "
           } else {
              System.write("L----")
              indent = indent + "|    "
           }
           var sColor = (root.color == 1) ? "RED" :"BLACK"
           System.print(root.data.toString + "(" + sColor + ")")
           printHelper_(root.left, indent, false)
           printHelper_(root.right, indent, true)
        }
    }

    /* public methods */

    // check if a node with a given key already exists in the tree
    hasKey(key) {
        var node = searchTreeHelper_(_root, key)
        return node.data == key
    }

    // pre-order traversal
    // node.left subtree.right subtree
    preOrder() { preOrderHelper_(_root) }

    // in-order traversal
    // left subtree . node . right subtree
    inOrder() { inOrderHelper_(_root) }

    // post-order traversal
    // left subtree . right Subtree . node
    postOrder() { postOrderHelper_(_root) }

    // search the tree for the key k
    // and return the corresponding node
    searchTree(k) { searchTreeHelper_(_root, k) }

    // find the node with the minimum key
    minimum(node) {
        while (node.left != _tnull) node = node.left
        return node
    }

    // find the node with the maximum key
    maximum(node) {
        while (node.right != _tnull) node = node.right
        return node
    }

    // find the successor of a given node
    successor(x) {
        // if the right subtree is not _tnull
        // the successor is the leftmost node in the
        // right subtree
        if (x.right != _tnull) return minimum(x.right)

        // else it is the lowest ancestor of x whose
        // right child is also an ancestor of x.
        var y = x.parent
        while (y != _tnull && x == y.right) {
            x = y
            y = y.parent
        }
        return y
    }

    // find the predecessor of a given node
    predecessor(x) {
        // if the left subtree is not _tnull
        // the predecessor is the rightmost node in the
        // left subtree
        if (x.left != _tnull) return maximum(x.left)

        // else it is the lowest ancestor of x whose
        // left child is also an ancestor of x.
        var y = x.parent
        while (y != _tnull && x == y.left) {
            x = y
            y = y.parent
        }
    }

    // rotate left at node x
    leftRotate(x) {
        var y = x.right
        x.right = y.left
        if (y.left != _tnull) y.left.parent = x
        y.parent = x.parent
        if (!x.parent) {
            _root = y
        } else if (x == x.parent.left) {
            x.parent.left = y
        } else {
            x.parent.right = y
        }
        y.left = x
        x.parent = y
    }

    // rotate right at node x
    rightRotate(x) {
        var y = x.left
        x.left = y.right
        if (y.right != _tnull) y.right.parent = x
        y.parent = x.parent
        if (!x.parent) {
            _root = y
        } else if (x == x.parent.right) {
            x.parent.right = y
        } else {
            x.parent.left = y
        }
        y.right = x
        x.parent = y
    }

    // insert the key to the tree in its appropriate position
    // and fix the tree
    insert(key) {
        // ordinary binary search insertion
        var y = null
        var x = _root
        while (x != _tnull) {
            y = x
            if (key < x.data) {
                x = x.left
            } else {
                x = x.right
            }
        }

        // if duplicates aren't allowed and the key has already been used, simply return
        if (!_allowDups && y && y.data == key) return

        // new node must be red
        var node = Node.new(key, null, _tnull, _tnull, 1) 

        // y is parent of x
        node.parent = y
        if (!y) {
            _root = node
        } else if (node.data < y.data) {
            y.left = node
        } else {
            y.right = node
        }

        // if new node is a root node, simply return
        if (!node.parent) {
            node.color = 0
            return
        }

        // if the grandparent is null, simply return
        if (!node.parent.parent) return

        // fix the tree
        fixInsert_(node)
    }

    // delete the node from the tree
    deleteNode(data) {
        deleteNodeHelper_(_root, data)
    }

    // print the tree structure on the terminal
    prettyPrint() {
        printHelper_(_root, "", true)
    }
}

var bst = RBTree.new(false) // disallow duplicates
var rand = Random.new()
// permit random keys from 1 to 99 say
var candidates = (1..99).toList
// shuffle them and insert the first 30 in the RBT
rand.shuffle(candidates)
var insertions = candidates[0..29]
for (k in insertions) bst.insert(k)
System.print("State of the tree after inserting the 30 keys:")
System.print(insertions)
System.print()
bst.prettyPrint()

// now delete 15 of them at random and redisplay the tree
rand.shuffle(insertions)
var deletions = insertions[0..14]
for (k in deletions) bst.deleteNode(k)
System.print("\nState of the tree after deleting the 15 keys:")
System.print(deletions)
System.print()
bst.prettyPrint()
Output:

Sample output.

State of the tree after inserting the 30 keys:
[85, 57, 51, 50, 86, 39, 95, 49, 44, 48, 21, 83, 11, 59, 88, 66, 4, 40, 24, 82, 63, 22, 37, 32, 91, 74, 28, 75, 62, 81]

R----50(BLACK)
     L----39(RED)
     |    L----21(BLACK)
     |    |    L----11(BLACK)
     |    |    |    L----4(RED)
     |    |    R----24(RED)
     |    |         L----22(BLACK)
     |    |         R----32(BLACK)
     |    |              L----28(RED)
     |    |              R----37(RED)
     |    R----44(BLACK)
     |         L----40(BLACK)
     |         R----49(BLACK)
     |              L----48(RED)
     R----83(RED)
          L----66(BLACK)
          |    L----57(RED)
          |    |    L----51(BLACK)
          |    |    R----62(BLACK)
          |    |         L----59(RED)
          |    |         R----63(RED)
          |    R----75(RED)
          |         L----74(BLACK)
          |         R----82(BLACK)
          |              L----81(RED)
          R----86(BLACK)
               L----85(BLACK)
               R----91(BLACK)
                    L----88(RED)
                    R----95(RED)

State of the tree after deleting the 15 keys:
[32, 40, 57, 63, 66, 75, 86, 59, 83, 51, 24, 62, 82, 39, 37]

R----50(BLACK)
     L----28(RED)
     |    L----21(BLACK)
     |    |    L----11(BLACK)
     |    |    |    L----4(RED)
     |    |    R----22(BLACK)
     |    R----48(BLACK)
     |         L----44(BLACK)
     |         R----49(BLACK)
     R----85(BLACK)
          L----81(BLACK)
          |    L----74(RED)
          R----91(RED)
               L----88(BLACK)
               R----95(BLACK)