tree/simple_tree.hpp

• View this file on GitHub
• Last update: 2024-10-30 04:43:18+09:00
• Include: #include "tree/simple_tree.hpp"

Depends on

• data_structure/csr.hpp

Verified with

• test/tree/FrequencyTableofTreeDistance.test.cpp

Code

#pragma once

#include <iostream>
#include "../data_structure/csr.hpp"

namespace noya2 {

struct simple_tree {
    internal::csr<int> g;
    simple_tree () {}
    simple_tree (int _n) : g(_n, (_n - 1)*2) {
        if (_n == 1){
            g.build();
        }
    }
    void add_edge(int u, int v){
        g.add(u, v);
        int id = g.add(v, u);
        if (id + 1 == (g.n - 1)*2) g.build();
    }
    void input(int indexed = 1){
        for (int i = 0; i < g.n - 1; i++){
            int u, v; cin >> u >> v;
            u -= indexed, v -= indexed;
            add_edge(u, v);
        }
    }
    void input_parents(int indexed = 1){
        for (int i = 0; i < g.n - 1; i++){
            int v; cin >> v;
            v -= indexed;
            add_edge(i + 1, v);
        }
    }
    const auto operator[](int v) const {
        return g[v];
    }
    auto operator[](int v){
        return g[v];
    }
    size_t size() const {
        return g.size();
    }
};

} // namespace noya2

#line 2 "tree/simple_tree.hpp"

#include <iostream>
#line 2 "data_structure/csr.hpp"

#include<vector>
#include<ranges>
#include<cassert>
#include<utility>

namespace noya2::internal {

template<class E>
struct csr {
    csr () {}
    csr (int _n) : n(_n) {}
    csr (int _n, int m) : n(_n){
        start.reserve(m);
        elist.reserve(m);
    }
    // ACL style constructor (do not have to call build)
    csr (int _n, const std::vector<std::pair<int,E>> &idx_elem) : n(_n), start(_n + 2), elist(idx_elem.size()) {
        for (auto &[i, e] : idx_elem){
            start[i + 2]++;
        }
        for (int i = 1; i < n; i++){
            start[i + 2] += start[i + 1];
        }
        for (auto &[i, e] : idx_elem){
            elist[start[i + 1]++] = e;
        }
        prepared = true;
    }
    int add(int idx, E elem){
        int eid = start.size();
        start.emplace_back(idx);
        elist.emplace_back(elem);
        return eid;
    }
    void build(){
        if (prepared) return ;
        int m = start.size();
        std::vector<E> nelist(m);
        std::vector<int> nstart(n + 2, 0);
        for (int i = 0; i < m; i++){
            nstart[start[i] + 2]++;
        }
        for (int i = 1; i < n; i++){
            nstart[i + 2] += nstart[i + 1];
        }
        for (int i = 0; i < m; i++){
            nelist[nstart[start[i] + 1]++] = elist[i];
        }
        swap(elist,nelist);
        swap(start,nstart);
        prepared = true;
    }
    const auto operator[](int idx) const {
        return std::ranges::subrange(elist.begin()+start[idx],elist.begin()+start[idx+1]);
    }
    auto operator[](int idx){
        return std::ranges::subrange(elist.begin()+start[idx],elist.begin()+start[idx+1]);
    }
    const auto operator()(int idx, int l, int r) const {
        return std::ranges::subrange(elist.begin()+start[idx]+l,elist.begin()+start[idx]+r);
    }
    auto operator()(int idx, int l, int r){
        return std::ranges::subrange(elist.begin()+start[idx]+l,elist.begin()+start[idx]+r);
    }
    size_t size() const {
        return n;
    }
    int n;
    std::vector<int> start;
    std::vector<E> elist;
    bool prepared = false;
};

} // namespace noya2::internal
#line 5 "tree/simple_tree.hpp"

namespace noya2 {

struct simple_tree {
    internal::csr<int> g;
    simple_tree () {}
    simple_tree (int _n) : g(_n, (_n - 1)*2) {
        if (_n == 1){
            g.build();
        }
    }
    void add_edge(int u, int v){
        g.add(u, v);
        int id = g.add(v, u);
        if (id + 1 == (g.n - 1)*2) g.build();
    }
    void input(int indexed = 1){
        for (int i = 0; i < g.n - 1; i++){
            int u, v; cin >> u >> v;
            u -= indexed, v -= indexed;
            add_edge(u, v);
        }
    }
    void input_parents(int indexed = 1){
        for (int i = 0; i < g.n - 1; i++){
            int v; cin >> v;
            v -= indexed;
            add_edge(i + 1, v);
        }
    }
    const auto operator[](int v) const {
        return g[v];
    }
    auto operator[](int v){
        return g[v];
    }
    size_t size() const {
        return g.size();
    }
};

} // namespace noya2

Back to top page