BinaryRadixTree.h

Go to the documentation of this file.

 
#ifndef PBAT_COMMON_BINARYRADIXTREE_H
#define PBAT_COMMON_BINARYRADIXTREE_H
 
#include "pbat/Aliases.h"
#include "pbat/common/Stack.h"
 
#include <algorithm>
#include <bit>
 
namespace pbat::common {

template <class TIndex = Index>
class BinaryRadixTree
{
  public:
    BinaryRadixTree() = default;
    template <class TDerived>
    BinaryRadixTree(Eigen::DenseBase<TDerived> const& codes, bool bStoreParent = false);

    template <class TDerived>
    void Construct(Eigen::DenseBase<TDerived> const& codes, bool bStoreParent = false);
    TIndex Left(TIndex i) const { return mChild(0, i); }
    TIndex Right(TIndex i) const { return mChild(1, i); }
    TIndex Parent(TIndex i) const { return mParent(i); }
    TIndex InternalNodeCount() const { return mChild.cols(); }
    TIndex LeafCount() const { return mParent.size() - InternalNodeCount(); }
    bool IsLeaf(TIndex i) const { return i >= InternalNodeCount(); }
    constexpr TIndex Root() const { return 0; }
    TIndex CodeIndex(TIndex leaf) const { return leaf - InternalNodeCount(); }
    TIndex LeafIndex(TIndex codeIdx) const { return codeIdx + InternalNodeCount(); }
    auto Children() const { return mChild; }
    auto Left() const { return mChild.row(0); }
    auto Right() const { return mChild.row(1); }
    auto Parent() const { return mParent; }
    bool HasParentRelationship() const { return mParent.size() > 0; }
 
  private:
    Eigen::Matrix<TIndex, 2, Eigen::Dynamic>
        mChild; 
    Eigen::Vector<TIndex, Eigen::Dynamic> mParent; 
};
 
template <class TIndex>
template <class TDerived>
inline BinaryRadixTree<TIndex>::BinaryRadixTree(
    Eigen::DenseBase<TDerived> const& codes,
    bool bStoreParent)
{
    Construct(codes.derived(), bStoreParent);
}
 
template <class TIndex>
template <class TDerived>
inline void
BinaryRadixTree<TIndex>::Construct(Eigen::DenseBase<TDerived> const& codes, bool bStoreParent)
{
    using CodeType = typename TDerived::Scalar;
    static_assert(
        std::is_integral_v<CodeType> and std::is_unsigned_v<CodeType> and
            not std::is_same_v<CodeType, bool>,
        "Codes must be of unsigned integral type");
    auto constexpr nBits = sizeof(CodeType) * 8;
    static_assert(nBits <= 64, "CodeType must have at most 64 bits");
 
    std::uint64_t constexpr msb = 0b1ULL << (nBits - 1);
    TIndex const nLeaves        = codes.size();
    TIndex const nInternal      = nLeaves - 1;
    mChild.resize(2, nInternal);
    mParent.resize(bStoreParent * (nLeaves + nInternal));
 
    struct Node
    {
        TIndex begin;
        TIndex end;
    };
    auto const fCommonPrefixLength = [](CodeType ci, CodeType cj) {
        return std::countl_zero(ci ^ cj);
    };
    pbat::common::Stack<Node, 64> stack{};
    stack.Push({0, nLeaves - 1});
    // Top-down construction over internal nodes
    while (not stack.IsEmpty())
    {
        Node const node = stack.Pop();
        // Compute range [first, last] of codes covered by the node.
        // If the node is a left child, its range is reversed (i.e. end -> begin).
        // If the node is a right child, its range is not reversed (i.e. begin -> end).
        bool const bReversed = node.begin > node.end;
        auto const first     = (not bReversed) * node.begin + bReversed * node.end;
        auto const last      = (not bReversed) * node.end + bReversed * node.begin;
        // Find the split position
        TIndex split      = first;
        auto const cfirst = codes(first);
        auto const clast  = codes(last);
        if (cfirst == clast)
        {
            split += (last - first + 1) / 2;
        }
        else
        {
            auto const mask  = msb >> fCommonPrefixLength(cfirst, clast);
            auto const begin = codes.begin() + first;
            auto const end   = codes.begin() + last + 1;
            auto const upper = std::upper_bound(begin, end, cfirst, [&](CodeType ci, CodeType cj) {
                return (mask & ci) < (mask & cj);
            });
            split += std::distance(begin, upper);
        }
        // The left and right child ranges are split as [first, split-1] and [split,last],
        // respectively.
        TIndex lc = split - 1;
        TIndex rc = split;
        // Ranges of size 1 indicate leaf nodes. We offset leaf indices by the number of internal
        // nodes.
        bool const bIsLeftLeaf  = (lc == first);
        bool const bIsRightLeaf = (rc == last);
        lc += bIsLeftLeaf * nInternal;
        rc += bIsRightLeaf * nInternal;
        // Set parent-child relationships
        mChild(0, node.begin) = lc;
        mChild(1, node.begin) = rc;
        // Only recurse into internal nodes
        if (not bIsLeftLeaf)
            stack.Push({lc, first});
        if (not bIsRightLeaf)
            stack.Push({rc, last});
    }
    // Store parent relationships if requested
    if (bStoreParent)
    {
        mParent(0) = 0;
        for (auto i = 0; i < nInternal; ++i)
        {
            mParent(Left(i))  = i;
            mParent(Right(i)) = i;
        }
    }
}
 
} // namespace pbat::common
 
#endif // PBAT_COMMON_BINARYRADIXTREE_H