I recently developed a simple treap data structure in C#. In the future, I may have the class extend IDictionary
. For now, I only expose a public accessor and a delete method.
public TValue this[TKey key]
{
get { return Get(root, key).Value; }
set { root = Set(root, new Node(key, value)); }
}
The data structure is generic and maintains the key-value pairs in an inner Node
class. Internally, the tree is organized both by the binary search requirement on Node.Key
and the heap invariant requirement on Node.Priority
.
public class BinaryTreap<TKey, TValue> where TKey : IComparable<TKey>
{
[DebuggerDisplay("Key={Key}, Priority={Priority}")]
/// <summary>Represents a Node in a BinaryTreap</summary>
class Node
{
private static Random random = new Random();
public readonly int Priority;
public readonly TKey Key;
public readonly TValue Value;
public Node Left, Right;
public Node(TKey key, TValue value,
Node left = null, Node right = null)
{
Priority = random.Next(); Key = key;
Value = value; Left = left; Right = right;
}
}
/// <summary>Tree root, organized as min-heap/BST</summary>
private Node root;
}
The methods Set
and Delete
are implemented similar to the standard binary search tree functions for adding/removing from the tree (these can be found on wikipedia). The only difference is that now an insertion or a deletion might trigger rotating the tree to ensure the heap invariant.
/// <summary>
/// Sets/Replaces node located at child.key in root
/// and rebalances the tree if necessary
/// </summary>
private Node Set(Node root, Node child)
{
if (root == null) return child;
int cmp = child.Key.CompareTo(root.Key);
if (cmp == 0)
return child;
else if (cmp > 0)
{
root.Right = Set(root.Right, child);
if (root.Right.Priority < root.Priority)
root = LRotate(root);
}
else if (cmp < 0)
{
root.Left = Set(root.Left, child);
if (root.Left.Priority < root.Priority)
root = RRotate(root);
}
return root;
}
/// <summary>
/// Deletes key from root by demoting it down to a leaf
/// and rebalances the tree if necessary
/// </summary>
private Node Delete(Node root, TKey key)
{
if (root == null) return null;
int cmp = key.CompareTo(root.Key);
if (cmp < 0) root.Left = Delete(root.Left, key);
if (cmp > 0) root.Right = Delete(root.Right, key);
if (cmp== 0)
{
if (root.Left != null && root.Right != null)
{
if (root.Left.Priority < root.Right.Priority)
root = RRotate(root);
else
root = LRotate(root);
}
else if (root.Left != null) root = RRotate(root);
else if (root.Right != null) root = LRotate(root);
else return null;
root = Delete(root, key);
}
return root;
}
I'm wondering if my tree rotation methods can be improved. Should I make the Nodes
in the tree immutable, and return a new rotated Node
instead?
/// <summary>
/// Promotes the right child to root
/// (i.e. rotate the "right child" left)
/// </summary>
private Node LRotate(Node root)
{
Node temp = root.Right;
root.Right = temp.Left;
temp.Left = root;
return temp;
}
/// <summary>
/// Promotes the left child to root
/// (i.e. rotate the "left child" right)
/// </summary>
private Node RRotate(Node root)
{
Node temp = root.Left;
root.Left = temp.Right;
temp.Right = root;
return temp;
}
Finally, my code provides validation. This is intended for testing purposes only, though I'm not sure how this should be exposed to a test project without making Node
internal and then exposing the project's internals. Maybe there's no better way.
/// <summary>Performs an in-order traversal from root</summary>
private IEnumerable<Node> InOrder(Node root)
{
if (root == null)
yield break;
foreach (Node child in InOrder(root.Left).Append(root).Concat(InOrder(root.Right)))
yield return child;
}
/// <summary>Validates the min-heap order for root</summary>
private bool IsHeap(Node root)
{
if (root == null)
return true;
if (root.Right != null && root.Right.Priority < root.Priority)
return false;
if (root.Left != null && root.Left.Priority < root.Priority)
return false;
return IsHeap(root.Left) && IsHeap(root.Right);
}
/// <summary>
/// Returns true if the heap order and bst order
/// properties are satisfied
/// </summary>
public bool AssertTreap()
{
return root == null ? true :
InOrder(root).SequenceEqual(InOrder(root).OrderBy(x => x.Key)) && IsHeap(root);
}
I've definedAppend
in an extension method
internal static class Extensions
{
public static IEnumerable<T> Append<T>(this IEnumerable<T> source, T item)
{
foreach (T x in source)
yield return x;
yield return item;
}
}
Example of use
Here's a simple test that can be run in the debugger.
class Program
{
public static void Main(string[] args)
{
var t = new BinaryTreap<int, int>();
for (int i = 0; i < 1000; i++)
t[i] = i;
bool valid = t.AssertTreap();
t.Delete(1);
try
{
int one = t[1];
}
catch
{
// doesn't exist
}
}
}
What I'd like reviewed
- Better (more functional) implementations of
RotateR
andRotateL
? - How should I expose
Node
to a test project? Right now I have my validation code in the class itself which doesn't seem right. I don't want to markNode
internal; is there a better way to expose the underlying data structure for testing purposes?
Edit
Here is the Get
method used by the indexer.
/// <summary>Retrieves Node located at key by binary search</summary>
private Node Get(Node root, TKey key)
{
if (root == null) return null;
int cmp = key.CompareTo(root.Key);
if (cmp < 0) return Get(root.Left, key);
if (cmp > 0) return Get(root.Right, key);
return root;
}
Here is the public wrapper for Delete
.
public void Delete(TKey key)
{
root = Delete(root, key);
}