/// <summary>Returns all sub-trees enumerated from left to right.</summary>
/// <returns>Enumerated sub-trees or empty if tree is empty.</returns>
public IEnumerable <ImMap <V> > Enumerate()
{
if (IsEmpty)
{
yield break;
}
var parents = new ImMap <V> [this.GetHeight()];
var node = this;
var parentCount = -1;
while (!node.IsEmpty || parentCount != -1)
{
if (!node.IsEmpty)
{
parents[++parentCount] = node;
node = node.GetLeft();
}
else
{
node = parents[parentCount--];
yield return(node);
node = node.GetRight();
}
}
}
internal Branch(int key, V value, ImMap <V> left, int rightHeight, ImMap <V> right) : base(key, value)
{
Left = left;
Right = right;
var leftHeight = left is Branch b ? b.Height : 1;
Height = leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
}
private static ImMap <V> BranchOrLeaf(int key, V value, ImMap <V> left, ImMap <V> right)
{
if (left == Empty && right == Empty)
{
return(new ImMap <V>(key, value));
}
return(new Branch(key, value, left, right));
}
///
public ImMap <V> AddOrUpdate(int key, V value)
{
var mapKey = Key;
var b = this as Branch;
if (b == null) // means the leaf node
{
// update the leaf
if (mapKey == key)
{
return(new ImMap <V>(key, value));
}
return(key < mapKey // search for node
? new Branch(mapKey, Value, new ImMap <V>(key, value), Empty, 2)
: new Branch(mapKey, Value, Empty, new ImMap <V>(key, value), 2));
}
// the empty branch node
var height = b.Height;
if (height == 0)
{
return(new ImMap <V>(key, value));
}
// update the branch key and value
var left = b.Left;
var right = b.Right;
if (mapKey == key)
{
return(new Branch(key, value, left, right, height));
}
if (key < mapKey)
{
left = left.AddOrUpdate(key, value);
}
else
{
right = right.AddOrUpdate(key, value);
}
// Now balance!!!
return(ImMap.Balance(mapKey, Value, left, right));
}
private static ImMap <V> Balance(int key, V value, ImMap <V> left, ImMap <V> right)
{
var delta = left.GetHeight() - right.GetHeight();
if (delta >= 2) // left is longer by 2, rotate left
{
var leftLeft = left.GetLeft();
var leftRight = left.GetRight();
if (leftRight.GetHeight() - leftLeft.GetHeight() == 1)
{
// double rotation:
// 5 => 5 => 4
// 2 6 4 6 2 5
// 1 4 2 3 1 3 6
// 3 1
return(new Branch(leftRight.Key, leftRight.Value,
BranchOrLeaf(left.Key, left.Value, leftLeft, leftRight.GetLeft()),
BranchOrLeaf(key, value, leftRight.GetRight(), right)));
}
// todo: do we need this?
// one rotation:
// 5 => 2
// 2 6 1 5
// 1 4 4 6
return(new Branch(left.Key, left.Value, leftLeft, BranchOrLeaf(key, value, leftRight, right)));
}
if (delta <= -2)
{
var rightLeft = right.GetLeft();
var rightRight = right.GetRight();
if (rightLeft.GetHeight() - rightRight.GetHeight() == 1)
{
return(new Branch(rightLeft.Key, rightLeft.Value,
BranchOrLeaf(key, value, left, rightLeft.GetLeft()),
BranchOrLeaf(right.Key, right.Value, rightLeft.GetRight(), rightRight)));
}
return(new Branch(right.Key, right.Value, BranchOrLeaf(key, value, left, rightLeft), rightRight));
}
return(new Branch(key, value, left, right));
}
internal Branch(int key, V value, ImMap <V> left, ImMap <V> right, int height) : base(key, value)
{
Left = left;
Right = right;
Height = height;
}
internal Branch(int key, V value, int leftHeight, ImMap <V> left, int rightHeight, ImMap <V> right) : base(key, value)
{
Left = left;
Right = right;
Height = leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
}
