public static ulong GetCountLeft(LeftLeaningRedBlackTree_WithEvents <int> tree, ulong[] counts, int val)
{
var cnt = 0ul;
foreach (var node in tree.Search(n => n.Key == val ? (bool?)null : n.Key < val ? false : true))
{
if (node.Key < val)
{
if (node.Left != null)
{
cnt += counts[node.Left.Key] + 1;
}
else if (node.Right == null)
{
cnt += 1;
}
}
else if (node.Key == val)
{
cnt += 1;
if (node.Left != null)
{
cnt += counts[node.Left.Key];
}
}
}
return(cnt);
}
public static ulong Solve_LLRBT(int n, int k, int[][] edges)
{
var adj = GraphHelper.ToAdjacencyList(edges, n);
var notHeads = new bool[n];
foreach (var edge in edges)
{
notHeads[edge[1]] = true;
}
var head = notHeads.FirstIndexOf(nh => !nh);
var count = 0ul;
var counts = new ulong[n];
var tree = new LeftLeaningRedBlackTree_WithEvents <int>(
(x, y) => x.CompareTo(y),
v => counts[v.Key] = 1ul + (v.Left == null ? 0ul : counts[v.Left.Key]) + (v.Right == null ? 0ul : counts[v.Right.Key]));
tree.Add(head);
foreach (var edgeInfo in GraphHelper.Traverse_Dfs_WithEvents(head, v => adj[v]))
{
var child = edgeInfo.Child;
if (edgeInfo.IsDownNotUp)
{
var countLeft = GetCountLeft(tree, counts, child - k - 1);
var countRight = GetCountLeft(tree, counts, child + k);
count += countRight - countLeft;
tree.Add(child);
}
else
{
tree.Remove(child);
}
}
return(count);
}
