k_ary_tree.hpp

// Copyright 2021 Arnaud Becheler    <abechele@umich.edu>
 
 
#pragma once
 
#include "detail/tree_traits.hpp"
#include "detail/visit.hpp"
 
#include <boost/range/adaptor/transformed.hpp>
#include <boost/range/iterator_range.hpp>
 
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/graphviz.hpp>
#include <boost/graph/isomorphism.hpp>
#include <boost/graph/topological_sort.hpp>
 
#include <boost/graph/connected_components.hpp>
#include <boost/graph/random.hpp>
#include <boost/graph/random_spanning_tree.hpp>
 
#include "detail/adl_resolution.hpp"
#include "detail/utils.hpp"
 
#include <algorithm>
#include <iostream>
 
#include <range/v3/all.hpp>
 
namespace quetzal::coalescence
{
using no_property = boost::no_property;
 
namespace detail
{
 
// Base class and common implementation
template <class VertexProperty, class EdgeProperty, class CRTP> class k_ary_tree_common
{
  protected:
    using base = tree_traits::model<tree_traits::out_edge_list_type, tree_traits::vertex_list_type,
                                    tree_traits::directed_type, VertexProperty, EdgeProperty>;
 
    base _graph;
 
  public:
    explicit k_ary_tree_common() : _graph()
    {
    }
 
    explicit k_ary_tree_common(size_t n) : _graph(n)
    {
    }
 
    using vertex_descriptor = typename base::vertex_descriptor;
 
    using edge_descriptor = typename base::edge_descriptor;
 
    using vertex_property = VertexProperty;
 
    using edge_property = EdgeProperty;
 
    using degree_size_type = typename base::degree_size_type;
 
    using traversal_category = typename base::traversal_category;
 
    using directed_category = typename base::directed_category;
 
    using in_edge_iterator = typename base::in_edge_iterator;
 
    using out_edge_iterator = typename base::out_edge_iterator;
 
    using edge_parallel_category = typename base::edge_parallel_category;
 
 
    bool is_isomorphic(k_ary_tree_common const &other) const
    {
        return boost::isomorphism(_graph, other._graph);
    }
 
    vertex_descriptor find_root_from(vertex_descriptor u) const
    {
        return has_predecessor(u) ? find_root_from(predecessor(u)) : u;
    }
 
    degree_size_type degree(vertex_descriptor u) const
    {
        return boost::degree(u, _graph);
    }
 
    degree_size_type in_degree(vertex_descriptor u) const
    {
        return has_predecessor(u);
    }
 
    degree_size_type out_degree(vertex_descriptor v) const
    {
        return boost::out_degree(v, _graph);
    }
 
    bool has_predecessor(vertex_descriptor u) const
    {
        auto range = boost::in_edges(u, _graph);
        return std::distance(range.first, range.second) > 0;
    }
 
    vertex_descriptor predecessor(vertex_descriptor u) const
    {
        auto range = boost::in_edges(u, _graph);
        assert(std::distance(range.first, range.second) == 1);
        return source(*range.first);
    }
 
    bool has_successors(vertex_descriptor u) const
    {
        auto range = boost::out_edges(u, _graph);
        assert(std::distance(range.first, range.second) != 1);
        return std::distance(range.first, range.second) >= 2;
    }
 
    auto successors(vertex_descriptor u) const
    {
        return boost::make_iterator_range(out_edges(u, _graph)) |
               boost::adaptors::transformed([this](auto it) { return target(it); });
    }
 
    std::optional<edge_descriptor> edge(vertex_descriptor u, vertex_descriptor v) const
    {
        auto p = boost::edge(u, v, _graph);
        return p.second ? std::make_optional(p.first) : std::nullopt;
    }
 
    vertex_descriptor source(edge_descriptor e) const
    {
        return boost::source(e, _graph);
    }
 
    vertex_descriptor target(edge_descriptor e) const
    {
        return boost::target(e, _graph);
    }
 
 
 
    std::pair<in_edge_iterator, in_edge_iterator> in_edges(vertex_descriptor u) const
    {
        using detail::adl_resolution::in_edges;
        return in_edges(u, _graph);
    }
 
 
 
    template <typename DFSTreeVisitor> void depth_first_search(vertex_descriptor start, DFSTreeVisitor &vis)
    {
        // convert tree visitor to graph visitor
        struct VisWrap : boost::default_dfs_visitor
        {
            k_ary_tree_common &_tree;
            DFSTreeVisitor &_tree_visitor;
            VisWrap(k_ary_tree_common &tree, DFSTreeVisitor &v) : _tree(tree), _tree_visitor(v)
            {
            }
            void discover_vertex(vertex_descriptor v, const base &g)
            {
                _tree_visitor(boost::visit::pre, v);
            }
            void tree_edge(const edge_descriptor &e, const base &g)
            {
                _tree_visitor(boost::visit::in, _tree.target(e));
            }
            void finish_vertex(vertex_descriptor v, const base &g)
            {
                _tree_visitor(boost::visit::post, v);
            }
        } graph_visitor(*this, vis);
        std::vector<boost::default_color_type> colors(boost::num_vertices(_graph));
        boost::iterator_property_map color_map(colors.begin(), boost::get(boost::vertex_index, _graph));
        return boost::depth_first_search(_graph, boost::visitor(graph_visitor).color_map(color_map).root_vertex(start));
    }
 
    template <typename DFSTreeVisitor> void depth_first_search(vertex_descriptor start, DFSTreeVisitor &vis) const
    {
        // convert tree visitor to graph visitor
        struct VisWrap : boost::default_dfs_visitor
        {
            const k_ary_tree_common &_tree;
            DFSTreeVisitor &_tree_visitor;
            VisWrap(const k_ary_tree_common &tree, DFSTreeVisitor &v) : _tree(tree), _tree_visitor(v)
            {
            }
            void discover_vertex(vertex_descriptor v, const base &g)
            {
                _tree_visitor(boost::visit::pre, v);
            }
            void tree_edge(const edge_descriptor &e, const base &g)
            {
                _tree_visitor(boost::visit::in, _tree.target(e));
            }
            void finish_vertex(vertex_descriptor v, const base &g)
            {
                _tree_visitor(boost::visit::post, v);
            }
        } graph_visitor(*this, vis);
        std::vector<boost::default_color_type> colors(boost::num_vertices(_graph));
        boost::iterator_property_map color_map(colors.begin(), boost::get(boost::vertex_index, _graph));
        return boost::depth_first_search(_graph, boost::visitor(graph_visitor).color_map(color_map).root_vertex(start));
    }
 
 
 
    void to_graphviz(std::ostream &out) const
    {
        using namespace boost;
        return boost::write_graphviz(out, *this, boost::make_label_writer(boost::get(vertex_bundle, *this)),
                                     boost::make_label_writer(boost::get(edge_bundle, *this)));
    }
 
 
    static constexpr inline vertex_descriptor null_vertex()
    {
        return base::null_vertex();
    }
 
  private:
    template <typename Generator>
    static auto update_tree(CRTP &tree, std::vector<vertex_descriptor> leaves, Generator &rng)
    {
        using tree_type = CRTP;
 
        if (leaves.size() == 1)
        {
            return leaves.front();
        }
        else
        {
            std::uniform_int_distribution<> distrib(1, leaves.size() - 1);
            int split = distrib(rng);
 
            std::vector<vertex_descriptor> left(leaves.begin(), leaves.begin() + split);
            std::vector<vertex_descriptor> right(leaves.begin() + split, leaves.end());
 
            auto parent = boost::add_vertex(tree._graph);
            if constexpr (std::is_same_v<EdgeProperty, boost::no_property>)
                tree.add_edges(parent, {update_tree(tree, left, rng), update_tree(tree, right, rng)});
            else
                tree.add_edges(parent, {std::make_tuple(update_tree(tree, left, rng), EdgeProperty()),
                                        std::make_tuple(update_tree(tree, right, rng), EdgeProperty())});
            return parent;
        }
    }
 
  public:
    template <typename Generator>
    static std::tuple<CRTP, vertex_descriptor> make_random_tree(int n_leaves, int k_max, Generator &rng)
    {
        CRTP tree;
        std::vector<vertex_descriptor> leaves(n_leaves);
        std::transform(leaves.cbegin(), leaves.cend(), leaves.begin(), [&tree](auto) { return tree.add_vertex(); });
        vertex_descriptor root = update_tree(tree, leaves, rng);
        return std::tuple(std::move(tree), root);
    }
 
}; // end class k_ary_tree_common
 
} // end namespace detail
 
template <class VertexProperty, class EdgeProperty> class k_ary_tree
{
};
 
template <>
class k_ary_tree<no_property, no_property>
    : public detail::k_ary_tree_common<no_property, no_property, k_ary_tree<no_property, no_property>>
{
    using base = detail::k_ary_tree_common<no_property, no_property, k_ary_tree<no_property, no_property>>;
 
  public:
 
    explicit k_ary_tree() : base()
    {
    }
 
    explicit k_ary_tree(size_t n) : base(n)
    {
    }
 
 
    vertex_descriptor add_vertex()
    {
        return boost::add_vertex(_graph);
    }
 
    std::vector<edge_descriptor> add_edges(vertex_descriptor parent, std::vector<vertex_descriptor> children)
    {
        // std::cout << children.size() << std::endl;
        assert(children.size() != 0);
        // for(auto const& c : children){ std::cout << parent << " -> " << c << std::endl; }
 
        assert(children.size() > 1);
        for (auto const &c : children)
        {
            assert(parent != c);
        }
        std::vector<edge_descriptor> edges(children.size());
        std::transform(children.cbegin(), children.cend(), edges.begin(),
                       [parent, this](auto c) { return boost::add_edge(parent, c, this->_graph).first; });
        return edges;
    }
 
 
}; // end specialization k_ary_tree<boost::no::property, boost::no_property>
 
template <class VertexProperty>
    requires(!std::is_same_v<VertexProperty, no_property>)
class k_ary_tree<VertexProperty, no_property>
    : public detail::k_ary_tree_common<VertexProperty, no_property, k_ary_tree<VertexProperty, no_property>>
{
  private:
    using base = detail::k_ary_tree_common<VertexProperty, no_property, k_ary_tree<VertexProperty, no_property>>;
 
  public:
    using edge_descriptor = typename base::edge_descriptor;
 
    using vertex_descriptor = typename base::vertex_descriptor;
 
 
    explicit k_ary_tree() : base()
    {
    }
 
    explicit k_ary_tree(size_t n) : base(n)
    {
    }
 
 
 
    auto add_vertex(const VertexProperty &p)
    {
        return boost::add_vertex(p, this->_graph);
    }
 
    std::vector<edge_descriptor> add_edges(vertex_descriptor parent, std::vector<vertex_descriptor> children)
    {
        assert(children.size() > 1);
        for (auto const &c : children)
        {
            assert(parent != c);
        }
 
        std::vector<edge_descriptor> edges(children.size());
        std::transform(children.cbegin(), children.cend(), edges.begin(),
                       [parent, this](auto c) { return boost::add_edge(parent, c, this->_graph).first; });
        return edges;
    }
 
 
 
    const VertexProperty &operator[](vertex_descriptor v) const
    {
        return this->_graph[v];
    }
 
    VertexProperty &operator[](vertex_descriptor v)
    {
        return this->_graph[v];
    }
 
 
}; // end specialization k_ary_tree<Vertex, boost::no_property>
 
template <class EdgeProperty>
    requires(!std::is_same_v<EdgeProperty, no_property>)
class k_ary_tree<no_property, EdgeProperty>
    : public detail::k_ary_tree_common<no_property, EdgeProperty, k_ary_tree<no_property, EdgeProperty>>
{
  private:
    using base = detail::k_ary_tree_common<no_property, EdgeProperty, k_ary_tree<no_property, EdgeProperty>>;
 
  public:
    using edge_descriptor = typename base::edge_descriptor;
 
    using vertex_descriptor = typename base::vertex_descriptor;
 
 
    explicit k_ary_tree() : base()
    {
    }
 
    explicit k_ary_tree(size_t n) : base(n)
    {
    }
 
 
 
    vertex_descriptor add_vertex()
    {
        return boost::add_vertex(this->_graph);
    }
 
    std::vector<edge_descriptor> add_edges(vertex_descriptor parent,
                                           const std::vector<std::pair<vertex_descriptor, EdgeProperty>> &children)
    {
        assert(children.size() > 1);
        for (auto const &[c, p] : children)
        {
            assert(parent != c);
        }
 
        std::vector<edge_descriptor> edges(children.size());
        std::transform(children.cbegin(), children.cend(), edges.begin(), [parent, this](const auto &c) {
            return boost::add_edge(parent, c.first, c.second, this->_graph).first;
        });
        return edges;
    }
 
 
 
    const EdgeProperty &operator[](const edge_descriptor &edge) const
    {
        return this->_graph[edge];
    }
 
    EdgeProperty &operator[](const edge_descriptor &edge)
    {
        return this->_graph[edge];
    }
 
 
}; // end specialization k_ary_tree<boost::no_property, Edge>
 
template <class VertexProperty, class EdgeProperty>
    requires(!std::is_same_v<VertexProperty, no_property> && !std::is_same_v<EdgeProperty, no_property>)
class k_ary_tree<VertexProperty, EdgeProperty>
    : public detail::k_ary_tree_common<VertexProperty, EdgeProperty, k_ary_tree<no_property, EdgeProperty>>
{
  private:
    using base = detail::k_ary_tree_common<VertexProperty, EdgeProperty, k_ary_tree<no_property, EdgeProperty>>;
 
  public:
    using edge_descriptor = typename base::edge_descriptor;
 
    using vertex_descriptor = typename base::vertex_descriptor;
 
 
 
    explicit k_ary_tree() : base()
    {
    }
 
    explicit k_ary_tree(size_t n) : base(n)
    {
    }
 
 
 
    vertex_descriptor add_vertex(const VertexProperty &p)
    {
        return boost::add_vertex(p, this->_graph);
    }
 
    std::vector<edge_descriptor> add_edges(vertex_descriptor parent,
                                           std::vector<std::tuple<vertex_descriptor, EdgeProperty>> children)
    {
        std::cout << children.size() << std::endl;
        assert(children.size() > 1);
        for (auto const &[c, p] : children)
        {
            assert(parent != c);
        }
 
        std::vector<edge_descriptor> edges(children.size());
        std::transform(children.cbegin(), children.cend(), edges.begin(), [parent, this](const auto &c) {
            return boost::add_edge(parent, std::get<0>(c), std::get<1>(c), this->_graph).first;
        });
        return edges;
    }
 
 
 
    const VertexProperty &operator[](vertex_descriptor v) const
    {
        return this->_graph[v];
    }
 
    VertexProperty &operator[](vertex_descriptor v)
    {
        return this->_graph[v];
    }
 
    const EdgeProperty &operator[](const edge_descriptor &edge) const
    {
        return this->_graph[edge];
    }
 
    EdgeProperty &operator[](const edge_descriptor &edge)
    {
        return this->_graph[edge];
    }
 
};
 
namespace detail
{
template <typename Vertex, typename Edge, typename Generator>
auto update_tree(k_ary_tree<Vertex, Edge> &tree,
                 std::vector<typename k_ary_tree<Vertex, Edge>::vertex_descriptor> leaves, Generator &rng)
{
    using tree_type = k_ary_tree<Vertex, Edge>;
    using vertex_descriptor = typename tree_type::vertex_descriptor;
 
    if (leaves.size() == 1)
    {
        return leaves.front();
    }
    else
    {
        std::uniform_int_distribution<> distrib(1, leaves.size() - 1);
        int split = distrib(rng);
 
        std::vector<vertex_descriptor> left(leaves.begin(), leaves.begin() + split);
        std::vector<vertex_descriptor> right(leaves.begin() + split, leaves.end());
 
        vertex_descriptor parent;
 
        if constexpr (std::is_same_v<Vertex, boost::no_property>)
            parent = tree.add_vertex();
        else
            parent = tree.add_vertex(Vertex());
 
        if constexpr (std::is_same_v<Edge, boost::no_property>)
            tree.add_edges(parent, {update_tree(tree, left, rng), update_tree(tree, right, rng)});
        else
            tree.add_edges(parent, {{update_tree(tree, left, rng), Edge()}, {update_tree(tree, right, rng), Edge()}});
        return parent;
    }
}
} // end namespace detail
 
template <typename Vertex = boost::no_property, typename Edge = boost::no_property, typename Generator>
std::pair<quetzal::coalescence::k_ary_tree<Vertex, Edge>,
          typename quetzal::coalescence::k_ary_tree<Vertex, Edge>::vertex_descriptor>
get_random_k_ary_tree(int n_leaves, Generator &rng)
{
    using tree_type = quetzal::coalescence::k_ary_tree<Vertex, Edge>;
    using vertex_descriptor = typename tree_type::vertex_descriptor;
 
    tree_type tree;
 
    std::vector<vertex_descriptor> leaves(n_leaves);
 
    std::transform(leaves.cbegin(), leaves.cend(), leaves.begin(), [&tree](auto) {
        if constexpr (std::is_same_v<Vertex, boost::no_property>)
            return tree.add_vertex();
        else
            return tree.add_vertex(Vertex());
    });
 
    vertex_descriptor root = detail::update_tree(tree, leaves, rng);
    return std::make_pair(std::move(tree), root);
}
 
} // end namespace quetzal::coalescence