Древесная Анархия [FPC]

type
  _tRBTNodeColor = (RBT_Red, RBT_Black);

  // Associative array. Based on red-black tree.
  generic gMap<_Key_, _Item_> = object(tRWObject)
  type
    tItemProc = procedure(var a: _Item_);
  private
  type
    _pRBTNode = ^_tRBTNode;
    _tRBTNode = record
      left, right, parent: _pRBTNode;
      color: _tRBTNodeColor;
      key: _Key_;
      data: _Item_;
    end;
  var
    _root: _pRBTNode;
    _count: sint;
    function _FindNode(const key: _Key_): _pRBTNode;
    procedure _RotateLeft(x: _pRBTNode);
    procedure _RotateRight(x: _pRBTNode);
    function _AddNode(const key: _Key_; const item: _Item_): _pRBTNode;
    procedure _DeleteNode(z: _pRBTNode);
    procedure _ForEach(x: _pRBTNode; proc: tItemProc);
    procedure _Free(x: _pRBTNode);
  protected
    function _NullItem: _Item_; virtual; abstract;
    procedure _Finalize(var key: _Key_; var item: _Item_); virtual;
    function _Compare(const a, b: _Key_): tValueRelationship; virtual; abstract;
  public
    constructor Init;
    destructor Done; virtual;
    procedure Add(const key: _Key_; const item: _Item_);
    function Find(const key: _Key_): _Item_;
    procedure Remove(const key: _Key_);
    procedure Clear;
    procedure ForEach(proc: tItemProc);
    property Count: sint read _count;
  end;

////////////////////////////////////////////////////////////////////

  function gMap._FindNode(const key: _Key_): _pRBTNode;
  var
    cn: _pRBTNode;
  begin
    cn := _root;
    while Assigned(cn) do
      case _Compare(key, cn^.key) of
        value_Equals: break;
        value_Greater: cn := cn^.right;
        value_Less: cn := cn^.left;
      end;
    result := cn;
  end;

  procedure gMap._RotateLeft(x: _pRBTNode);
  var
    y: _pRBTNode;
  begin
    y := x^.right;
    x^.right := y^.left;
    if Assigned(y^.left) then y^.left^.parent := x;
    y^.parent := x^.parent;
    if x = _root then _root := y else
      if x = x^.parent^.left then
        x^.parent^.left := y
      else
        x^.parent^.right := y;
    y^.left := x;
    x^.parent := y;
  end;

  procedure gMap._RotateRight(x: _pRBTNode);
  var
    y: _pRBTNode;
  begin
    y := x^.left;
    x^.left := y^.right;
    if Assigned(y^.right) then y^.right^.parent := x;
    y^.parent := x^.parent;
    if x = _root then _root := y else
      if x = x^.parent^.right then
        x^.parent^.right := y
      else
        x^.parent^.left := y;
    y^.right := x;
    x^.parent := y;
  end;

  function gMap._AddNode(const key: _Key_; const item: _Item_): _pRBTNode;
  var
    x, y, z, zpp: _pRBTNode;
  begin
    y := nil;
    x := _root;
    while Assigned(x) do
    begin
      y := x;
      case _Compare(key, x^.key) of
        value_Equals:
          begin
            _Finalize(x^.key, x^.data);
            x^.key := key;
            x^.data := item;
            exit(x); // already exists
          end;
        value_Less: x := x^.left;
        value_Greater: x := x^.right;
      end;
    end;

    inc(_count);
    new(z);
    z^.key := key;
    z^.data := item;
    z^.left := nil;
    z^.right := nil;
    z^.color := RBT_Red;
    z^.parent := y;
    result := z;

    if Assigned(y) then
      case _Compare(key, y^.key) of
        value_Less: y^.left := z;
        else
          y^.right := z;
      end
    else
      _root := z;

    // rebalance
    while (z <> _root) and (z^.parent^.color = RBT_Red) do
    begin
      zpp := z^.parent^.parent;
      if z^.parent = zpp^.left then
      begin
        y := zpp^.right;
        if Assigned(y) and (y^.color = RBT_Red) then
        begin
          z^.parent^.color := RBT_Black;
          y^.color := RBT_Black;
          zpp^.color := RBT_Red;
          z := zpp;
        end else
        begin
          if z = z^.parent^.right then
          begin
            z := z^.parent;
            _RotateLeft(z);
          end;
          z^.parent^.color := RBT_Black;
          zpp^.color := RBT_Red;
          _RotateRight(zpp);
        end;
      end else
      begin
        y := zpp^.left;
        if Assigned(y) and (y^.color = RBT_Red) then
        begin
          z^.parent^.color := RBT_Black;
          y^.color := RBT_Black;
          zpp^.color := RBT_Red;
          z := zpp;
        end else
        begin
          if z = z^.parent^.left then
          begin
            z := z^.parent;
            _RotateRight(z);
          end;
          z^.parent^.color := RBT_Black;
          zpp^.color := RBT_Red;
          _RotateLeft(zpp);
        end;
      end;
    end;
    _root^.color := RBT_Black;
  end;

  procedure gMap._DeleteNode(z: _pRBTNode);
  var
    w, x, y, x_parent: _pRBTNode;
    tmpcol: _tRBTNodeColor;
  begin
    y := z;
    x := nil;
    x_parent := nil;

    if Assigned(y^.left) then
    begin
      if Assigned(y^.right) then
      begin
        y := y^.right;
        while Assigned(y^.left) do y := y^.left;
        x := y^.right;
      end else
        x := y^.left;
    end else
      x := y^.right;

    if y <> z then
    begin
      // relink y in place of z. y is z's successor
      z^.left^.parent := y;
      y^.left := z^.left;
      if y <> z^.right then
      begin
        x_parent := y^.parent;
        if Assigned(x) then x^.parent := y^.parent;
        y^.parent^.left := x; // y must be a child of left
        y^.right := z^.right;
        z^.right^.parent := y;
      end else
        x_parent := y;
      if _root = z then
        _root := y
      else
        if z^.parent^.left = z then
          z^.parent^.left := y
        else
          z^.parent^.right := y;
      y^.parent := z^.parent;
      tmpcol := y^.color;
      y^.color := z^.color;
      z^.color := tmpcol;
      y := z;
    end else
    begin // y = z
      x_parent := y^.parent;
      if Assigned(x)  then
        x^.parent := y^.parent;
      if _root = z then
        _root := x
      else
        if z^.parent^.left = z then
          z^.parent^.left := x
        else
          z^.parent^.right := x;
    end;

    // rebalance
    if y^.color = RBT_Black then
    begin
      while (x <> _root) and ((not Assigned(x)) or (x^.color = RBT_Black)) do
      begin
        if x = x_parent^.left then
        begin
          w := x_parent^.right;
          if w^.color = RBT_Red then
          begin
            w^.color := RBT_Black;
            x_parent^.color := RBT_Red;
            _RotateLeft(x_parent);
            w := x_parent^.right;
          end;
          if ((not Assigned(w^.left)) or (w^.left^.color = RBT_Black)) and
             ((not Assigned(w^.right)) or (w^.right^.color = RBT_Black)) then
          begin
            w^.color := RBT_Red;
            x := x_parent;
            x_parent := x_parent^.parent;
          end else
          begin
            if (not Assigned(w^.right)) or (w^.right^.color = RBT_Black) then
            begin
              w^.left^.color := RBT_Black;
              w^.color := RBT_Red;
              _RotateRight(w);
              w := x_parent^.right;
            end;
            w^.color := x_parent^.color;
            x_parent^.color := RBT_Black;
            if Assigned(w^.right)  then w^.right^.color := RBT_Black;
            _RotateLeft(x_parent);
            x := _root; // break;
          end;
        end else
        begin
          // mirror of above code
          w := x_parent^.left;
          if (w^.color = RBT_Red)  then
          begin
            w^.color := RBT_Black;
            x_parent^.color := RBT_Red;
            _RotateRight(x_parent);
            w := x_parent^.left;
          end;
          if ((not Assigned(w^.right)) or (w^.right^.color = RBT_Black)) and
             ((not Assigned(w^.left)) or (w^.left^.color = RBT_Black)) then
          begin
            w^.color := RBT_Red;
            x := x_parent;
            x_parent := x_parent^.parent;
          end else
          begin
            if (not Assigned(w^.left)) or (w^.left^.color = RBT_Black) then
            begin
              w^.right^.color := RBT_Black;
              w^.color := RBT_Red;
              _RotateLeft(w);
              w := x_parent^.left;
            end;
            w^.color := x_parent^.color;
            x_parent^.color := RBT_Black;
            if Assigned(w^.left) then w^.left^.color := RBT_Black;
            _RotateRight(x_parent);
            x := _root; // break;
          end;
        end;
      end;
      if Assigned(x) then x^.color := RBT_Black;
    end;

    _Finalize(y^.key, y^.data);
    dispose(y);
    dec(_count);
  end;

  procedure gMap._ForEach(x: _pRBTNode; proc: tItemProc);
  begin
    if Assigned(x) then
    begin
      _ForEach(x^.left, proc);
      _ForEach(x^.right, proc);
      proc(x^.data);
    end;
  end;

  procedure gMap._Free(x: _pRBTNode);
  begin
    if Assigned(x) then
    begin
      _Free(x^.left);
      _Free(x^.right);
      _Finalize(x^.key, x^.data);
      dispose(x);
    end;
  end;

  procedure gMap._Finalize(var key: _Key_; var item: _Item_);
  begin
    Assert((@key = @key) and (@item = @item));
  end;

  constructor gMap.Init;
  begin
    inherited Init;
    _count := 0;
  end;

  destructor gMap.Done;
  begin
    Clear;
    inherited Done;
  end;

  procedure gMap.Add(const key: _Key_; const item: _Item_);
  begin
    _AddNode(key, item);
  end;

  procedure gMap.Remove(const key: _Key_);
  var
    n: _pRBTNode;
  begin
    n := _FindNode(key);
    if Assigned(n) then _DeleteNode(n);
  end;

  procedure gMap.Clear;
  begin
    _Free(_root);
    _root := nil;
  end;

  function gMap.Find(const key: _Key_): _Item_;
  var
    n: _pRBTNode;
  begin
    n := _FindNode(key);
    if Assigned(n) then
      result := n^.data
    else
      result := _nullItem;
  end;

  procedure gMap.ForEach(proc: tItemProc);
  begin
    _ForEach(_root, proc);
  end;