tree_contour.hpp¶
Centroid-decomposition distance-range updates and queries on a static tree.
Verified by vertex_add_range_contour_sum_on_tree, vertex_get_range_contour_add_on_tree.
在静态树上按到某点的距离区间做更新或查询；适合“距离在 [l,r) 内的所有点”这类操作。
Implementation¶
#ifndef NOYA_TREE_CONTOUR_HPP
#define NOYA_TREE_CONTOUR_HPP 1
/// @complexity Time: O(n log^2 n) construction and O(log^2 n) per update or
/// query. Space: O(n log n).
#include <algorithm>
#include <cassert>
#include <tuple>
#include <utility>
#include <vector>
namespace noya {
/// @brief Maintain values indexed by tree distance. Every vertex is recorded
/// in all centroid-ancestor buckets; querying adds all buckets and subtracts
/// the bucket of the branch containing the query vertex. Fenwick trees over
/// distances support both point-add/range-sum and the dual range-add/point-get.
class tree_contour {
struct fenwick {
std::vector<long long> data;
fenwick() = default;
explicit fenwick(int size) : data(size + 1) {}
int size() const { return int(data.size()) - 1; }
void add(int index, long long value) {
assert(0 <= index && index < size());
for (index++; index < int(data.size()); index += index & -index) {
data[index] += value;
}
}
long long prefix_sum(int count) const {
count = std::clamp(count, 0, size());
long long result = 0;
for (; count > 0; count -= count & -count) {
result += data[count];
}
return result;
}
};
struct bucket {
int distance_count = 0;
fenwick point_values;
fenwick range_tags;
void reset(int count) {
distance_count = count;
point_values = fenwick(count);
range_tags = fenwick(count + 1);
}
void point_add(int distance, long long value) {
point_values.add(distance, value);
}
long long point_prefix(int distance_bound) const {
return point_values.prefix_sum(distance_bound);
}
void range_add(int left, int right, long long value) {
left = std::clamp(left, 0, distance_count);
right = std::clamp(right, 0, distance_count);
if (left >= right) {
return;
}
range_tags.add(left, value);
range_tags.add(right, -value);
}
long long point_tag(int distance) const {
assert(0 <= distance && distance < distance_count);
return range_tags.prefix_sum(distance + 1);
}
};
struct ancestor_entry {
int centroid = -1;
int distance = 0;
int branch = -1;
};
int vertex_count_ = 0;
std::vector<std::vector<int>> graph_;
std::vector<char> removed_;
std::vector<int> component_parent_;
std::vector<int> subtree_size_;
std::vector<std::vector<ancestor_entry>> ancestors_;
std::vector<bucket> all_;
std::vector<std::vector<bucket>> branches_;
std::vector<long long> initial_values_;
int find_centroid(int start) {
std::vector<int> order;
order.push_back(start);
component_parent_[start] = -1;
for (int index = 0; index < int(order.size()); index++) {
int vertex = order[index];
for (int next : graph_[vertex]) {
if (removed_[next] || next == component_parent_[vertex]) {
continue;
}
component_parent_[next] = vertex;
order.push_back(next);
}
}
for (int index = int(order.size()) - 1; index >= 0; index--) {
int vertex = order[index];
subtree_size_[vertex] = 1;
for (int next : graph_[vertex]) {
if (!removed_[next] && component_parent_[next] == vertex) {
subtree_size_[vertex] += subtree_size_[next];
}
}
}
int component_size = int(order.size());
int centroid = start;
int best_largest_part = component_size;
for (int vertex : order) {
int largest_part = component_size - subtree_size_[vertex];
for (int next : graph_[vertex]) {
if (!removed_[next] && component_parent_[next] == vertex) {
largest_part = std::max(largest_part, subtree_size_[next]);
}
}
if (largest_part < best_largest_part) {
best_largest_part = largest_part;
centroid = vertex;
}
}
return centroid;
}
void decompose(int start) {
int centroid = find_centroid(start);
removed_[centroid] = true;
ancestors_[centroid].push_back({centroid, 0, -1});
int maximum_distance = 0;
for (int neighbor : graph_[centroid]) {
if (removed_[neighbor]) {
continue;
}
int branch = int(branches_[centroid].size());
int branch_maximum_distance = 0;
std::vector<std::tuple<int, int, int>> stack = {
{neighbor, centroid, 1}};
while (!stack.empty()) {
auto [vertex, parent, distance] = stack.back();
stack.pop_back();
ancestors_[vertex].push_back({centroid, distance, branch});
branch_maximum_distance =
std::max(branch_maximum_distance, distance);
for (int next : graph_[vertex]) {
if (!removed_[next] && next != parent) {
stack.push_back({next, vertex, distance + 1});
}
}
}
branches_[centroid].push_back({});
branches_[centroid].back().reset(branch_maximum_distance + 1);
maximum_distance =
std::max(maximum_distance, branch_maximum_distance);
}
all_[centroid].reset(maximum_distance + 1);
for (int neighbor : graph_[centroid]) {
if (!removed_[neighbor]) {
decompose(neighbor);
}
}
}
long long prefix_contour_sum(int vertex, int distance_bound) const {
long long result = 0;
for (const ancestor_entry &entry : ancestors_[vertex]) {
int remaining = distance_bound - entry.distance;
result += all_[entry.centroid].point_prefix(remaining);
if (entry.branch != -1) {
result -= branches_[entry.centroid][entry.branch].point_prefix(remaining);
}
}
return result;
}
public:
explicit tree_contour(const std::vector<std::vector<int>> &graph,
const std::vector<long long> &initial_values = {})
: vertex_count_(int(graph.size())), graph_(graph),
removed_(vertex_count_), component_parent_(vertex_count_),
subtree_size_(vertex_count_), ancestors_(vertex_count_),
all_(vertex_count_), branches_(vertex_count_),
initial_values_(initial_values.empty()
? std::vector<long long>(vertex_count_)
: initial_values) {
assert(int(initial_values_.size()) == vertex_count_);
if (vertex_count_ == 0) {
return;
}
decompose(0);
for (int vertex = 0; vertex < vertex_count_; vertex++) {
point_add(vertex, initial_values_[vertex]);
}
}
/// @brief Add value to one vertex for later contour-sum queries.
void point_add(int vertex, long long value) {
assert(0 <= vertex && vertex < vertex_count_);
for (const ancestor_entry &entry : ancestors_[vertex]) {
all_[entry.centroid].point_add(entry.distance, value);
if (entry.branch != -1) {
branches_[entry.centroid][entry.branch].point_add(entry.distance,
value);
}
}
}
/// @brief Sum values at vertices whose distance from vertex is in [l,r).
long long range_sum(int vertex, int left, int right) const {
assert(0 <= vertex && vertex < vertex_count_ && left <= right);
return prefix_contour_sum(vertex, right) -
prefix_contour_sum(vertex, left);
}
/// @brief Add value to vertices whose distance from vertex is in [l,r).
void range_add(int vertex, int left, int right, long long value) {
assert(0 <= vertex && vertex < vertex_count_ && left <= right);
for (const ancestor_entry &entry : ancestors_[vertex]) {
int shifted_left = left - entry.distance;
int shifted_right = right - entry.distance;
all_[entry.centroid].range_add(shifted_left, shifted_right, value);
if (entry.branch != -1) {
branches_[entry.centroid][entry.branch].range_add(
shifted_left, shifted_right, value);
}
}
}
/// @brief Return the initial value plus every contour-range addition that
/// contains this vertex.
long long point_get(int vertex) const {
assert(0 <= vertex && vertex < vertex_count_);
long long result = initial_values_[vertex];
for (const ancestor_entry &entry : ancestors_[vertex]) {
result += all_[entry.centroid].point_tag(entry.distance);
if (entry.branch != -1) {
result -=
branches_[entry.centroid][entry.branch].point_tag(entry.distance);
}
}
return result;
}
};
} // namespace noya
#endif // NOYA_TREE_CONTOUR_HPP
#include <algorithm>
#include <cassert>
#include <tuple>
#include <utility>
#include <vector>
/// @complexity Time: O(n log^2 n) construction and O(log^2 n) per update or
/// query. Space: O(n log n).
namespace noya {
/// @brief Maintain values indexed by tree distance. Every vertex is recorded
/// in all centroid-ancestor buckets; querying adds all buckets and subtracts
/// the bucket of the branch containing the query vertex. Fenwick trees over
/// distances support both point-add/range-sum and the dual range-add/point-get.
class tree_contour {
struct fenwick {
std::vector<long long> data;
fenwick() = default;
explicit fenwick(int size) : data(size + 1) {}
int size() const { return int(data.size()) - 1; }
void add(int index, long long value) {
assert(0 <= index && index < size());
for (index++; index < int(data.size()); index += index & -index) {
data[index] += value;
}
}
long long prefix_sum(int count) const {
count = std::clamp(count, 0, size());
long long result = 0;
for (; count > 0; count -= count & -count) {
result += data[count];
}
return result;
}
};
struct bucket {
int distance_count = 0;
fenwick point_values;
fenwick range_tags;
void reset(int count) {
distance_count = count;
point_values = fenwick(count);
range_tags = fenwick(count + 1);
}
void point_add(int distance, long long value) {
point_values.add(distance, value);
}
long long point_prefix(int distance_bound) const {
return point_values.prefix_sum(distance_bound);
}
void range_add(int left, int right, long long value) {
left = std::clamp(left, 0, distance_count);
right = std::clamp(right, 0, distance_count);
if (left >= right) {
return;
}
range_tags.add(left, value);
range_tags.add(right, -value);
}
long long point_tag(int distance) const {
assert(0 <= distance && distance < distance_count);
return range_tags.prefix_sum(distance + 1);
}
};
struct ancestor_entry {
int centroid = -1;
int distance = 0;
int branch = -1;
};
int vertex_count_ = 0;
std::vector<std::vector<int>> graph_;
std::vector<char> removed_;
std::vector<int> component_parent_;
std::vector<int> subtree_size_;
std::vector<std::vector<ancestor_entry>> ancestors_;
std::vector<bucket> all_;
std::vector<std::vector<bucket>> branches_;
std::vector<long long> initial_values_;
int find_centroid(int start) {
std::vector<int> order;
order.push_back(start);
component_parent_[start] = -1;
for (int index = 0; index < int(order.size()); index++) {
int vertex = order[index];
for (int next : graph_[vertex]) {
if (removed_[next] || next == component_parent_[vertex]) {
continue;
}
component_parent_[next] = vertex;
order.push_back(next);
}
}
for (int index = int(order.size()) - 1; index >= 0; index--) {
int vertex = order[index];
subtree_size_[vertex] = 1;
for (int next : graph_[vertex]) {
if (!removed_[next] && component_parent_[next] == vertex) {
subtree_size_[vertex] += subtree_size_[next];
}
}
}
int component_size = int(order.size());
int centroid = start;
int best_largest_part = component_size;
for (int vertex : order) {
int largest_part = component_size - subtree_size_[vertex];
for (int next : graph_[vertex]) {
if (!removed_[next] && component_parent_[next] == vertex) {
largest_part = std::max(largest_part, subtree_size_[next]);
}
}
if (largest_part < best_largest_part) {
best_largest_part = largest_part;
centroid = vertex;
}
}
return centroid;
}
void decompose(int start) {
int centroid = find_centroid(start);
removed_[centroid] = true;
ancestors_[centroid].push_back({centroid, 0, -1});
int maximum_distance = 0;
for (int neighbor : graph_[centroid]) {
if (removed_[neighbor]) {
continue;
}
int branch = int(branches_[centroid].size());
int branch_maximum_distance = 0;
std::vector<std::tuple<int, int, int>> stack = {
{neighbor, centroid, 1}};
while (!stack.empty()) {
auto [vertex, parent, distance] = stack.back();
stack.pop_back();
ancestors_[vertex].push_back({centroid, distance, branch});
branch_maximum_distance =
std::max(branch_maximum_distance, distance);
for (int next : graph_[vertex]) {
if (!removed_[next] && next != parent) {
stack.push_back({next, vertex, distance + 1});
}
}
}
branches_[centroid].push_back({});
branches_[centroid].back().reset(branch_maximum_distance + 1);
maximum_distance =
std::max(maximum_distance, branch_maximum_distance);
}
all_[centroid].reset(maximum_distance + 1);
for (int neighbor : graph_[centroid]) {
if (!removed_[neighbor]) {
decompose(neighbor);
}
}
}
long long prefix_contour_sum(int vertex, int distance_bound) const {
long long result = 0;
for (const ancestor_entry &entry : ancestors_[vertex]) {
int remaining = distance_bound - entry.distance;
result += all_[entry.centroid].point_prefix(remaining);
if (entry.branch != -1) {
result -= branches_[entry.centroid][entry.branch].point_prefix(remaining);
}
}
return result;
}
public:
explicit tree_contour(const std::vector<std::vector<int>> &graph,
const std::vector<long long> &initial_values = {})
: vertex_count_(int(graph.size())), graph_(graph),
removed_(vertex_count_), component_parent_(vertex_count_),
subtree_size_(vertex_count_), ancestors_(vertex_count_),
all_(vertex_count_), branches_(vertex_count_),
initial_values_(initial_values.empty()
? std::vector<long long>(vertex_count_)
: initial_values) {
assert(int(initial_values_.size()) == vertex_count_);
if (vertex_count_ == 0) {
return;
}
decompose(0);
for (int vertex = 0; vertex < vertex_count_; vertex++) {
point_add(vertex, initial_values_[vertex]);
}
}
/// @brief Add value to one vertex for later contour-sum queries.
void point_add(int vertex, long long value) {
assert(0 <= vertex && vertex < vertex_count_);
for (const ancestor_entry &entry : ancestors_[vertex]) {
all_[entry.centroid].point_add(entry.distance, value);
if (entry.branch != -1) {
branches_[entry.centroid][entry.branch].point_add(entry.distance,
value);
}
}
}
/// @brief Sum values at vertices whose distance from vertex is in [l,r).
long long range_sum(int vertex, int left, int right) const {
assert(0 <= vertex && vertex < vertex_count_ && left <= right);
return prefix_contour_sum(vertex, right) -
prefix_contour_sum(vertex, left);
}
/// @brief Add value to vertices whose distance from vertex is in [l,r).
void range_add(int vertex, int left, int right, long long value) {
assert(0 <= vertex && vertex < vertex_count_ && left <= right);
for (const ancestor_entry &entry : ancestors_[vertex]) {
int shifted_left = left - entry.distance;
int shifted_right = right - entry.distance;
all_[entry.centroid].range_add(shifted_left, shifted_right, value);
if (entry.branch != -1) {
branches_[entry.centroid][entry.branch].range_add(
shifted_left, shifted_right, value);
}
}
}
/// @brief Return the initial value plus every contour-range addition that
/// contains this vertex.
long long point_get(int vertex) const {
assert(0 <= vertex && vertex < vertex_count_);
long long result = initial_values_[vertex];
for (const ancestor_entry &entry : ancestors_[vertex]) {
result += all_[entry.centroid].point_tag(entry.distance);
if (entry.branch != -1) {
result -=
branches_[entry.centroid][entry.branch].point_tag(entry.distance);
}
}
return result;
}
};
} // namespace noya