dynamic_tree_subtree.hpp

SECTIONGraph INCLUDEnoya/dynamic_tree_subtree.hpp

Link-Cut Tree with dynamic-edge subtree additions and sums.

Verified by dynamic_tree_subtree_add_subtree_sum, dynamic_tree_vertex_add_subtree_sum.

在可连边断边的森林中维护以指定根观察的子树加法与子树和。

Implementation

View on GitHub

SourceExpanded

#ifndef NOYA_DYNAMIC_TREE_SUBTREE_HPP
#define NOYA_DYNAMIC_TREE_SUBTREE_HPP 1

/// @complexity Time: Amortized O(log n) per link, cut, vertex update, or
/// directed-subtree operation. Space: O(n).

#include <cassert>
#include <cstdint>
#include <vector>

namespace noya {

/// @brief Link-Cut Tree augmented with virtual-child size and sum. Access
/// moves preferred children into or out of the virtual aggregate; each node
/// remembers how much of its represented parent's lazy addition it has
/// already received, so subtree additions remain valid across preferred-path
/// changes.
class dynamic_tree_subtree {
  struct node {
    node *parent = nullptr;
    node *child[2] = {nullptr, nullptr};
    long long value = 0;
    int size = 0;
    int virtual_size = 0;
    long long sum = 0;
    long long virtual_sum = 0;
    long long added = 0;
    long long parent_added_seen = 0;
    bool reversed = false;

    void apply_add(long long delta) {
      value += delta;
      sum += static_cast<long long>(size) * delta;
      virtual_sum += static_cast<long long>(virtual_size) * delta;
      added += delta;
    }

    void push(node *none) {
      if (parent != none) {
        apply_add(parent->added - parent_added_seen);
        parent_added_seen = parent->added;
      }
      if (reversed) {
        child[0]->reversed ^= child[0] != none;
        child[1]->reversed ^= child[1] != none;
        node *temporary = child[0];
        child[0] = child[1];
        child[1] = temporary;
        reversed = false;
      }
    }

    void pull(node *none) {
      if (child[0] != none) {
        child[0]->push(none);
      }
      if (child[1] != none) {
        child[1]->push(none);
      }
      size = 1 + child[0]->size + child[1]->size + virtual_size;
      sum = value + child[0]->sum + child[1]->sum + virtual_sum;
    }

    bool is_auxiliary_root(node *none) const {
      return parent == none ||
             (parent->child[0] != this && parent->child[1] != this);
    }

    void rotate(int direction, node *none) {
      node *new_root = child[direction ^ 1];
      node *middle = new_root->child[direction];
      if (middle != none) {
        middle->push(none);
      }
      child[direction ^ 1] = middle;
      if (middle != none) {
        middle->parent = this;
        middle->parent_added_seen = added;
      }
      new_root->child[direction] = this;
      new_root->parent = parent;
      new_root->parent_added_seen = parent->added;
      if (parent->child[0] == this) {
        parent->child[0] = new_root;
      }
      if (parent->child[1] == this) {
        parent->child[1] = new_root;
      }
      parent = new_root;
      parent_added_seen = new_root->added;
      pull(none);
      new_root->pull(none);
    }

    void splay(node *none) {
      while (!is_auxiliary_root(none)) {
        if (parent->is_auxiliary_root(none)) {
          parent->push(none);
          push(none);
          parent->rotate(parent->child[0] == this, none);
        } else {
          node *grandparent = parent->parent;
          grandparent->push(none);
          parent->push(none);
          push(none);
          bool parent_is_left = grandparent->child[0] == parent;
          bool vertex_is_left = parent->child[0] == this;
          if (parent_is_left == vertex_is_left) {
            grandparent->rotate(parent_is_left, none);
          }
          parent->rotate(vertex_is_left, none);
          if (parent_is_left != vertex_is_left) {
            grandparent->rotate(parent_is_left, none);
          }
        }
      }
      push(none);
    }

    void access(node *none) {
      node *previous = none;
      for (node *current = this; current != none; current = current->parent) {
        current->splay(none);
        if (current->child[1] != none) {
          current->child[1]->push(none);
          current->virtual_size += current->child[1]->size;
          current->virtual_sum += current->child[1]->sum;
        }
        current->child[1] = previous;
        if (previous != none) {
          previous->push(none);
          current->virtual_size -= previous->size;
          current->virtual_sum -= previous->sum;
        }
        current->pull(none);
        previous = current;
      }
      splay(none);
    }

    void make_root(node *none) {
      access(none);
      reversed = !reversed;
      push(none);
    }

    void link_to(node *new_parent, node *none) {
      access(none);
      new_parent->access(none);
      new_parent->child[1] = this;
      parent = new_parent;
      parent_added_seen = new_parent->added;
      new_parent->pull(none);
    }

    void cut_from_parent(node *none) {
      access(none);
      assert(child[0] != none);
      child[0]->push(none);
      child[0]->parent = none;
      child[0] = none;
      pull(none);
    }
  };

  node none_;
  std::vector<node> nodes_;

  node *none() { return &none_; }

  void check_vertex(int vertex) const {
    assert(0 <= vertex && vertex < int(nodes_.size()));
  }

public:
  explicit dynamic_tree_subtree(const std::vector<long long> &values)
      : nodes_(values.size()) {
    none_.parent = &none_;
    none_.child[0] = none_.child[1] = &none_;
    for (int vertex = 0; vertex < int(nodes_.size()); vertex++) {
      node &current = nodes_[vertex];
      current.parent = none();
      current.child[0] = current.child[1] = none();
      current.value = current.sum = values[vertex];
      current.size = 1;
    }
  }

  dynamic_tree_subtree(const dynamic_tree_subtree &) = delete;
  dynamic_tree_subtree &operator=(const dynamic_tree_subtree &) = delete;
  dynamic_tree_subtree(dynamic_tree_subtree &&) = delete;
  dynamic_tree_subtree &operator=(dynamic_tree_subtree &&) = delete;

  int size() const { return int(nodes_.size()); }

  /// @brief Link two different represented trees with an undirected edge.
  void link(int first, int second) {
    check_vertex(first);
    check_vertex(second);
    nodes_[first].make_root(none());
    nodes_[first].link_to(&nodes_[second], none());
  }

  /// @brief Remove an existing undirected edge.
  void cut(int first, int second) {
    check_vertex(first);
    check_vertex(second);
    nodes_[first].make_root(none());
    nodes_[second].cut_from_parent(none());
  }

  /// @brief Add delta to one vertex.
  void vertex_add(int vertex, long long delta) {
    check_vertex(vertex);
    nodes_[vertex].access(none());
    nodes_[vertex].value += delta;
    nodes_[vertex].pull(none());
  }

  /// @brief Add delta to the component on vertex's side after conceptually
  /// removing the adjacent edge (vertex,parent), then restore that edge.
  void subtree_add(int vertex, int parent, long long delta) {
    check_vertex(vertex);
    check_vertex(parent);
    nodes_[parent].make_root(none());
    nodes_[vertex].cut_from_parent(none());
    nodes_[vertex].apply_add(delta);
    nodes_[vertex].link_to(&nodes_[parent], none());
  }

  /// @brief Sum the component on vertex's side after conceptually removing
  /// the adjacent edge (vertex,parent), then restore that edge.
  long long subtree_sum(int vertex, int parent) {
    check_vertex(vertex);
    check_vertex(parent);
    nodes_[parent].make_root(none());
    nodes_[vertex].cut_from_parent(none());
    long long result = nodes_[vertex].sum;
    nodes_[vertex].link_to(&nodes_[parent], none());
    return result;
  }
};

} // namespace noya

#endif // NOYA_DYNAMIC_TREE_SUBTREE_HPP

#include <cassert>
#include <cstdint>
#include <vector>

/// @complexity Time: Amortized O(log n) per link, cut, vertex update, or
/// directed-subtree operation. Space: O(n).

namespace noya {

/// @brief Link-Cut Tree augmented with virtual-child size and sum. Access
/// moves preferred children into or out of the virtual aggregate; each node
/// remembers how much of its represented parent's lazy addition it has
/// already received, so subtree additions remain valid across preferred-path
/// changes.
class dynamic_tree_subtree {
  struct node {
    node *parent = nullptr;
    node *child[2] = {nullptr, nullptr};
    long long value = 0;
    int size = 0;
    int virtual_size = 0;
    long long sum = 0;
    long long virtual_sum = 0;
    long long added = 0;
    long long parent_added_seen = 0;
    bool reversed = false;

    void apply_add(long long delta) {
      value += delta;
      sum += static_cast<long long>(size) * delta;
      virtual_sum += static_cast<long long>(virtual_size) * delta;
      added += delta;
    }

    void push(node *none) {
      if (parent != none) {
        apply_add(parent->added - parent_added_seen);
        parent_added_seen = parent->added;
      }
      if (reversed) {
        child[0]->reversed ^= child[0] != none;
        child[1]->reversed ^= child[1] != none;
        node *temporary = child[0];
        child[0] = child[1];
        child[1] = temporary;
        reversed = false;
      }
    }

    void pull(node *none) {
      if (child[0] != none) {
        child[0]->push(none);
      }
      if (child[1] != none) {
        child[1]->push(none);
      }
      size = 1 + child[0]->size + child[1]->size + virtual_size;
      sum = value + child[0]->sum + child[1]->sum + virtual_sum;
    }

    bool is_auxiliary_root(node *none) const {
      return parent == none ||
             (parent->child[0] != this && parent->child[1] != this);
    }

    void rotate(int direction, node *none) {
      node *new_root = child[direction ^ 1];
      node *middle = new_root->child[direction];
      if (middle != none) {
        middle->push(none);
      }
      child[direction ^ 1] = middle;
      if (middle != none) {
        middle->parent = this;
        middle->parent_added_seen = added;
      }
      new_root->child[direction] = this;
      new_root->parent = parent;
      new_root->parent_added_seen = parent->added;
      if (parent->child[0] == this) {
        parent->child[0] = new_root;
      }
      if (parent->child[1] == this) {
        parent->child[1] = new_root;
      }
      parent = new_root;
      parent_added_seen = new_root->added;
      pull(none);
      new_root->pull(none);
    }

    void splay(node *none) {
      while (!is_auxiliary_root(none)) {
        if (parent->is_auxiliary_root(none)) {
          parent->push(none);
          push(none);
          parent->rotate(parent->child[0] == this, none);
        } else {
          node *grandparent = parent->parent;
          grandparent->push(none);
          parent->push(none);
          push(none);
          bool parent_is_left = grandparent->child[0] == parent;
          bool vertex_is_left = parent->child[0] == this;
          if (parent_is_left == vertex_is_left) {
            grandparent->rotate(parent_is_left, none);
          }
          parent->rotate(vertex_is_left, none);
          if (parent_is_left != vertex_is_left) {
            grandparent->rotate(parent_is_left, none);
          }
        }
      }
      push(none);
    }

    void access(node *none) {
      node *previous = none;
      for (node *current = this; current != none; current = current->parent) {
        current->splay(none);
        if (current->child[1] != none) {
          current->child[1]->push(none);
          current->virtual_size += current->child[1]->size;
          current->virtual_sum += current->child[1]->sum;
        }
        current->child[1] = previous;
        if (previous != none) {
          previous->push(none);
          current->virtual_size -= previous->size;
          current->virtual_sum -= previous->sum;
        }
        current->pull(none);
        previous = current;
      }
      splay(none);
    }

    void make_root(node *none) {
      access(none);
      reversed = !reversed;
      push(none);
    }

    void link_to(node *new_parent, node *none) {
      access(none);
      new_parent->access(none);
      new_parent->child[1] = this;
      parent = new_parent;
      parent_added_seen = new_parent->added;
      new_parent->pull(none);
    }

    void cut_from_parent(node *none) {
      access(none);
      assert(child[0] != none);
      child[0]->push(none);
      child[0]->parent = none;
      child[0] = none;
      pull(none);
    }
  };

  node none_;
  std::vector<node> nodes_;

  node *none() { return &none_; }

  void check_vertex(int vertex) const {
    assert(0 <= vertex && vertex < int(nodes_.size()));
  }

public:
  explicit dynamic_tree_subtree(const std::vector<long long> &values)
      : nodes_(values.size()) {
    none_.parent = &none_;
    none_.child[0] = none_.child[1] = &none_;
    for (int vertex = 0; vertex < int(nodes_.size()); vertex++) {
      node &current = nodes_[vertex];
      current.parent = none();
      current.child[0] = current.child[1] = none();
      current.value = current.sum = values[vertex];
      current.size = 1;
    }
  }

  dynamic_tree_subtree(const dynamic_tree_subtree &) = delete;
  dynamic_tree_subtree &operator=(const dynamic_tree_subtree &) = delete;
  dynamic_tree_subtree(dynamic_tree_subtree &&) = delete;
  dynamic_tree_subtree &operator=(dynamic_tree_subtree &&) = delete;

  int size() const { return int(nodes_.size()); }

  /// @brief Link two different represented trees with an undirected edge.
  void link(int first, int second) {
    check_vertex(first);
    check_vertex(second);
    nodes_[first].make_root(none());
    nodes_[first].link_to(&nodes_[second], none());
  }

  /// @brief Remove an existing undirected edge.
  void cut(int first, int second) {
    check_vertex(first);
    check_vertex(second);
    nodes_[first].make_root(none());
    nodes_[second].cut_from_parent(none());
  }

  /// @brief Add delta to one vertex.
  void vertex_add(int vertex, long long delta) {
    check_vertex(vertex);
    nodes_[vertex].access(none());
    nodes_[vertex].value += delta;
    nodes_[vertex].pull(none());
  }

  /// @brief Add delta to the component on vertex's side after conceptually
  /// removing the adjacent edge (vertex,parent), then restore that edge.
  void subtree_add(int vertex, int parent, long long delta) {
    check_vertex(vertex);
    check_vertex(parent);
    nodes_[parent].make_root(none());
    nodes_[vertex].cut_from_parent(none());
    nodes_[vertex].apply_add(delta);
    nodes_[vertex].link_to(&nodes_[parent], none());
  }

  /// @brief Sum the component on vertex's side after conceptually removing
  /// the adjacent edge (vertex,parent), then restore that edge.
  long long subtree_sum(int vertex, int parent) {
    check_vertex(vertex);
    check_vertex(parent);
    nodes_[parent].make_root(none());
    nodes_[vertex].cut_from_parent(none());
    long long result = nodes_[vertex].sum;
    nodes_[vertex].link_to(&nodes_[parent], none());
    return result;
  }
};

} // namespace noya