using System;
using System.Collections.Generic;
namespace ARMeilleure.Translation
{
///
/// An Augmented Interval Tree based off of the "TreeDictionary"'s Red-Black Tree. Allows fast overlap checking of ranges.
///
/// Key
/// Value
class IntervalTree where TK : IComparable
{
private const int ArrayGrowthSize = 32;
private const bool Black = true;
private const bool Red = false;
private IntervalTreeNode _root = null;
private int _count = 0;
public int Count => _count;
#region Public Methods
///
/// Gets the values of the interval whose key is .
///
/// Key of the node value to get
/// Value with the given
/// True if the key is on the dictionary, false otherwise
public bool TryGet(TK key, out TV value)
{
IntervalTreeNode node = GetNode(key);
if (node == null)
{
value = default;
return false;
}
value = node.Value;
return true;
}
///
/// Returns the start addresses of the intervals whose start and end keys overlap the given range.
///
/// Start of the range
/// End of the range
/// Overlaps array to place results in
/// Index to start writing results into the array. Defaults to 0
/// Number of intervals found
public int Get(TK start, TK end, ref TK[] overlaps, int overlapCount = 0)
{
GetKeys(_root, start, end, ref overlaps, ref overlapCount);
return overlapCount;
}
///
/// Adds a new interval into the tree whose start is , end is and value is .
///
/// Start of the range to add
/// End of the range to insert
/// Value to add
/// Optional factory used to create a new value if is already on the tree
/// is null
/// True if the value was added, false if the start key was already in the dictionary
public bool AddOrUpdate(TK start, TK end, TV value, Func updateFactoryCallback)
{
ArgumentNullException.ThrowIfNull(value);
return BSTInsert(start, end, value, updateFactoryCallback, out _);
}
///
/// Gets an existing or adds a new interval into the tree whose start is , end is and value is .
///
/// Start of the range to add
/// End of the range to insert
/// Value to add
/// is null
/// if is not yet on the tree, or the existing value otherwise
public TV GetOrAdd(TK start, TK end, TV value)
{
ArgumentNullException.ThrowIfNull(value);
BSTInsert(start, end, value, null, out IntervalTreeNode node);
return node.Value;
}
///
/// Removes a value from the tree, searching for it with .
///
/// Key of the node to remove
/// Number of deleted values
public int Remove(TK key)
{
int removed = Delete(key);
_count -= removed;
return removed;
}
///
/// Adds all the nodes in the dictionary into .
///
/// A list of all values sorted by Key Order
public List AsList()
{
List list = new();
AddToList(_root, list);
return list;
}
#endregion
#region Private Methods (BST)
///
/// Adds all values that are children of or contained within into , in Key Order.
///
/// The node to search for values within
/// The list to add values to
private void AddToList(IntervalTreeNode node, List list)
{
if (node == null)
{
return;
}
AddToList(node.Left, list);
list.Add(node.Value);
AddToList(node.Right, list);
}
///
/// Retrieve the node reference whose key is , or null if no such node exists.
///
/// Key of the node to get
/// is null
/// Node reference in the tree
private IntervalTreeNode GetNode(TK key)
{
ArgumentNullException.ThrowIfNull(key);
IntervalTreeNode node = _root;
while (node != null)
{
int cmp = key.CompareTo(node.Start);
if (cmp < 0)
{
node = node.Left;
}
else if (cmp > 0)
{
node = node.Right;
}
else
{
return node;
}
}
return null;
}
///
/// Retrieve all keys that overlap the given start and end keys.
///
/// Start of the range
/// End of the range
/// Overlaps array to place results in
/// Overlaps count to update
private void GetKeys(IntervalTreeNode node, TK start, TK end, ref TK[] overlaps, ref int overlapCount)
{
if (node == null || start.CompareTo(node.Max) >= 0)
{
return;
}
GetKeys(node.Left, start, end, ref overlaps, ref overlapCount);
bool endsOnRight = end.CompareTo(node.Start) > 0;
if (endsOnRight)
{
if (start.CompareTo(node.End) < 0)
{
if (overlaps.Length <= overlapCount)
{
Array.Resize(ref overlaps, overlapCount + ArrayGrowthSize);
}
overlaps[overlapCount++] = node.Start;
}
GetKeys(node.Right, start, end, ref overlaps, ref overlapCount);
}
}
///
/// Propagate an increase in max value starting at the given node, heading up the tree.
/// This should only be called if the max increases - not for rebalancing or removals.
///
/// The node to start propagating from
private static void PropagateIncrease(IntervalTreeNode node)
{
TK max = node.Max;
IntervalTreeNode ptr = node;
while ((ptr = ptr.Parent) != null)
{
if (max.CompareTo(ptr.Max) > 0)
{
ptr.Max = max;
}
else
{
break;
}
}
}
///
/// Propagate recalculating max value starting at the given node, heading up the tree.
/// This fully recalculates the max value from all children when there is potential for it to decrease.
///
/// The node to start propagating from
private static void PropagateFull(IntervalTreeNode node)
{
IntervalTreeNode ptr = node;
do
{
TK max = ptr.End;
if (ptr.Left != null && ptr.Left.Max.CompareTo(max) > 0)
{
max = ptr.Left.Max;
}
if (ptr.Right != null && ptr.Right.Max.CompareTo(max) > 0)
{
max = ptr.Right.Max;
}
ptr.Max = max;
} while ((ptr = ptr.Parent) != null);
}
///
/// Insertion Mechanism for the interval tree. Similar to a BST insert, with the start of the range as the key.
/// Iterates the tree starting from the root and inserts a new node where all children in the left subtree are less than , and all children in the right subtree are greater than .
/// Each node can contain multiple values, and has an end address which is the maximum of all those values.
/// Post insertion, the "max" value of the node and all parents are updated.
///
/// Start of the range to insert
/// End of the range to insert
/// Value to insert
/// Optional factory used to create a new value if is already on the tree
/// Node that was inserted or modified
/// True if was not yet on the tree, false otherwise
private bool BSTInsert(TK start, TK end, TV value, Func updateFactoryCallback, out IntervalTreeNode outNode)
{
IntervalTreeNode parent = null;
IntervalTreeNode node = _root;
while (node != null)
{
parent = node;
int cmp = start.CompareTo(node.Start);
if (cmp < 0)
{
node = node.Left;
}
else if (cmp > 0)
{
node = node.Right;
}
else
{
outNode = node;
if (updateFactoryCallback != null)
{
// Replace
node.Value = updateFactoryCallback(start, node.Value);
int endCmp = end.CompareTo(node.End);
if (endCmp > 0)
{
node.End = end;
if (end.CompareTo(node.Max) > 0)
{
node.Max = end;
PropagateIncrease(node);
RestoreBalanceAfterInsertion(node);
}
}
else if (endCmp < 0)
{
node.End = end;
PropagateFull(node);
}
}
return false;
}
}
IntervalTreeNode newNode = new(start, end, value, parent);
if (newNode.Parent == null)
{
_root = newNode;
}
else if (start.CompareTo(parent.Start) < 0)
{
parent.Left = newNode;
}
else
{
parent.Right = newNode;
}
PropagateIncrease(newNode);
_count++;
RestoreBalanceAfterInsertion(newNode);
outNode = newNode;
return true;
}
///
/// Removes the value from the dictionary after searching for it with .
///
/// Key to search for
/// Number of deleted values
private int Delete(TK key)
{
IntervalTreeNode nodeToDelete = GetNode(key);
if (nodeToDelete == null)
{
return 0;
}
IntervalTreeNode replacementNode;
if (LeftOf(nodeToDelete) == null || RightOf(nodeToDelete) == null)
{
replacementNode = nodeToDelete;
}
else
{
replacementNode = PredecessorOf(nodeToDelete);
}
IntervalTreeNode tmp = LeftOf(replacementNode) ?? RightOf(replacementNode);
if (tmp != null)
{
tmp.Parent = ParentOf(replacementNode);
}
if (ParentOf(replacementNode) == null)
{
_root = tmp;
}
else if (replacementNode == LeftOf(ParentOf(replacementNode)))
{
ParentOf(replacementNode).Left = tmp;
}
else
{
ParentOf(replacementNode).Right = tmp;
}
if (replacementNode != nodeToDelete)
{
nodeToDelete.Start = replacementNode.Start;
nodeToDelete.Value = replacementNode.Value;
nodeToDelete.End = replacementNode.End;
nodeToDelete.Max = replacementNode.Max;
}
PropagateFull(replacementNode);
if (tmp != null && ColorOf(replacementNode) == Black)
{
RestoreBalanceAfterRemoval(tmp);
}
return 1;
}
///
/// Returns the node with the largest key where is considered the root node.
///
/// Root Node
/// Node with the maximum key in the tree of
private static IntervalTreeNode Maximum(IntervalTreeNode node)
{
IntervalTreeNode tmp = node;
while (tmp.Right != null)
{
tmp = tmp.Right;
}
return tmp;
}
///
/// Finds the node whose key is immediately less than .
///
/// Node to find the predecessor of
/// Predecessor of
private static IntervalTreeNode PredecessorOf(IntervalTreeNode node)
{
if (node.Left != null)
{
return Maximum(node.Left);
}
IntervalTreeNode parent = node.Parent;
while (parent != null && node == parent.Left)
{
node = parent;
parent = parent.Parent;
}
return parent;
}
#endregion
#region Private Methods (RBL)
private void RestoreBalanceAfterRemoval(IntervalTreeNode balanceNode)
{
IntervalTreeNode ptr = balanceNode;
while (ptr != _root && ColorOf(ptr) == Black)
{
if (ptr == LeftOf(ParentOf(ptr)))
{
IntervalTreeNode sibling = RightOf(ParentOf(ptr));
if (ColorOf(sibling) == Red)
{
SetColor(sibling, Black);
SetColor(ParentOf(ptr), Red);
RotateLeft(ParentOf(ptr));
sibling = RightOf(ParentOf(ptr));
}
if (ColorOf(LeftOf(sibling)) == Black && ColorOf(RightOf(sibling)) == Black)
{
SetColor(sibling, Red);
ptr = ParentOf(ptr);
}
else
{
if (ColorOf(RightOf(sibling)) == Black)
{
SetColor(LeftOf(sibling), Black);
SetColor(sibling, Red);
RotateRight(sibling);
sibling = RightOf(ParentOf(ptr));
}
SetColor(sibling, ColorOf(ParentOf(ptr)));
SetColor(ParentOf(ptr), Black);
SetColor(RightOf(sibling), Black);
RotateLeft(ParentOf(ptr));
ptr = _root;
}
}
else
{
IntervalTreeNode sibling = LeftOf(ParentOf(ptr));
if (ColorOf(sibling) == Red)
{
SetColor(sibling, Black);
SetColor(ParentOf(ptr), Red);
RotateRight(ParentOf(ptr));
sibling = LeftOf(ParentOf(ptr));
}
if (ColorOf(RightOf(sibling)) == Black && ColorOf(LeftOf(sibling)) == Black)
{
SetColor(sibling, Red);
ptr = ParentOf(ptr);
}
else
{
if (ColorOf(LeftOf(sibling)) == Black)
{
SetColor(RightOf(sibling), Black);
SetColor(sibling, Red);
RotateLeft(sibling);
sibling = LeftOf(ParentOf(ptr));
}
SetColor(sibling, ColorOf(ParentOf(ptr)));
SetColor(ParentOf(ptr), Black);
SetColor(LeftOf(sibling), Black);
RotateRight(ParentOf(ptr));
ptr = _root;
}
}
}
SetColor(ptr, Black);
}
private void RestoreBalanceAfterInsertion(IntervalTreeNode balanceNode)
{
SetColor(balanceNode, Red);
while (balanceNode != null && balanceNode != _root && ColorOf(ParentOf(balanceNode)) == Red)
{
if (ParentOf(balanceNode) == LeftOf(ParentOf(ParentOf(balanceNode))))
{
IntervalTreeNode sibling = RightOf(ParentOf(ParentOf(balanceNode)));
if (ColorOf(sibling) == Red)
{
SetColor(ParentOf(balanceNode), Black);
SetColor(sibling, Black);
SetColor(ParentOf(ParentOf(balanceNode)), Red);
balanceNode = ParentOf(ParentOf(balanceNode));
}
else
{
if (balanceNode == RightOf(ParentOf(balanceNode)))
{
balanceNode = ParentOf(balanceNode);
RotateLeft(balanceNode);
}
SetColor(ParentOf(balanceNode), Black);
SetColor(ParentOf(ParentOf(balanceNode)), Red);
RotateRight(ParentOf(ParentOf(balanceNode)));
}
}
else
{
IntervalTreeNode sibling = LeftOf(ParentOf(ParentOf(balanceNode)));
if (ColorOf(sibling) == Red)
{
SetColor(ParentOf(balanceNode), Black);
SetColor(sibling, Black);
SetColor(ParentOf(ParentOf(balanceNode)), Red);
balanceNode = ParentOf(ParentOf(balanceNode));
}
else
{
if (balanceNode == LeftOf(ParentOf(balanceNode)))
{
balanceNode = ParentOf(balanceNode);
RotateRight(balanceNode);
}
SetColor(ParentOf(balanceNode), Black);
SetColor(ParentOf(ParentOf(balanceNode)), Red);
RotateLeft(ParentOf(ParentOf(balanceNode)));
}
}
}
SetColor(_root, Black);
}
private void RotateLeft(IntervalTreeNode node)
{
if (node != null)
{
IntervalTreeNode right = RightOf(node);
node.Right = LeftOf(right);
if (node.Right != null)
{
node.Right.Parent = node;
}
IntervalTreeNode nodeParent = ParentOf(node);
right.Parent = nodeParent;
if (nodeParent == null)
{
_root = right;
}
else if (node == LeftOf(nodeParent))
{
nodeParent.Left = right;
}
else
{
nodeParent.Right = right;
}
right.Left = node;
node.Parent = right;
PropagateFull(node);
}
}
private void RotateRight(IntervalTreeNode node)
{
if (node != null)
{
IntervalTreeNode left = LeftOf(node);
node.Left = RightOf(left);
if (node.Left != null)
{
node.Left.Parent = node;
}
IntervalTreeNode nodeParent = ParentOf(node);
left.Parent = nodeParent;
if (nodeParent == null)
{
_root = left;
}
else if (node == RightOf(nodeParent))
{
nodeParent.Right = left;
}
else
{
nodeParent.Left = left;
}
left.Right = node;
node.Parent = left;
PropagateFull(node);
}
}
#endregion
#region Safety-Methods
// These methods save memory by allowing us to forego sentinel nil nodes, as well as serve as protection against NullReferenceExceptions.
///
/// Returns the color of , or Black if it is null.
///
/// Node
/// The boolean color of , or black if null
private static bool ColorOf(IntervalTreeNode node)
{
return node == null || node.Color;
}
///
/// Sets the color of node to .
///
/// This method does nothing if is null.
///
/// Node to set the color of
/// Color (Boolean)
private static void SetColor(IntervalTreeNode node, bool color)
{
if (node != null)
{
node.Color = color;
}
}
///
/// This method returns the left node of , or null if is null.
///
/// Node to retrieve the left child from
/// Left child of
private static IntervalTreeNode LeftOf(IntervalTreeNode node)
{
return node?.Left;
}
///
/// This method returns the right node of , or null if is null.
///
/// Node to retrieve the right child from
/// Right child of
private static IntervalTreeNode RightOf(IntervalTreeNode node)
{
return node?.Right;
}
///
/// Returns the parent node of , or null if is null.
///
/// Node to retrieve the parent from
/// Parent of
private static IntervalTreeNode ParentOf(IntervalTreeNode node)
{
return node?.Parent;
}
#endregion
public bool ContainsKey(TK key)
{
return GetNode(key) != null;
}
public void Clear()
{
_root = null;
_count = 0;
}
}
///
/// Represents a node in the IntervalTree which contains start and end keys of type K, and a value of generic type V.
///
/// Key type of the node
/// Value type of the node
class IntervalTreeNode
{
public bool Color = true;
public IntervalTreeNode Left = null;
public IntervalTreeNode Right = null;
public IntervalTreeNode Parent = null;
///
/// The start of the range.
///
public TK Start;
///
/// The end of the range.
///
public TK End;
///
/// The maximum end value of this node and all its children.
///
public TK Max;
///
/// Value stored on this node.
///
public TV Value;
public IntervalTreeNode(TK start, TK end, TV value, IntervalTreeNode parent)
{
Start = start;
End = end;
Max = end;
Value = value;
Parent = parent;
}
}
}