Ежедневный бит (е) C ++ # 142, Общая задача на собеседовании: сумма расстояний до всех узлов.
Сегодня мы рассмотрим распространенную задачу интервью C++: сумму расстояний до всех узлов.
Учитывая дерево с n узлами, представленное в виде графа с использованием карты окрестностей, вычислите сумму расстояний до всех других узлов для каждого узла.
Идентификаторы узлов находятся в диапазоне [0,n).
Например, в приведенном выше дереве сумма расстояний для узлов 0 и 1 равна 4, а для узлов 2 и 3 — 6.
Прежде чем вы продолжите читать решение, я рекомендую вам попробовать решить его самостоятельно. Вот ссылка на Compiler Explorer с парой тестов: https://compiler-explorer.com/z/5nj1ejG5G.
Решение
Начнем с неоптимального решения. Рассмотрим поддерево, для корня которого мы уже вычислили сумму расстояний. Переход к родительскому узлу означает, что мы на один шаг дальше от всех узлов в этом поддереве.
Поэтому, если мы выполним прямой обход дерева в обратном порядке, мы сможем вычислить сумму расстояний для корневого узла, просто суммируя расстояния до дочерних поддеревьев плюс количество узлов в поддеревьях (поскольку это один шаг вперед будет вносить по одному на каждый узел).
Важно отметить, что это дает нам сумму расстояний только для корневого узла, поэтому мы должны повторить процесс для каждого узла (рассматривая каждый узел как корень дерева). Поскольку обход — это операция O(n), мы получаем временную сложность O(n*n).
Если вы некоторое время следили за этой серией, у вас может возникнуть подозрение, что мы делаем много повторяющейся работы и, возможно, есть возможности для улучшения.
Рассмотрим два соседних узла в дереве.
Когда мы вычислим их соответствующую сумму расстояний, формулы будут такими:
- сумма_расстояний(x) = сумма_поддеревьев(x) + сумма_поддеревьев(y) + количество_узлов(y)
- сумма_расстояний(y) = сумма_поддеревьев(y) + сумма_поддеревьев(x) + количество_узлов(x)
- сумма_расстояний(x) — сумма_расстояний(y) = количество_узлов(y) — количество_узлов(x)
Эта формула дает нам возможность вычислить ответ для ребенка из значения родителя.
- сумма_расстояний(дочерние) = сумма_расстояний(родительские) + количество_узлов(родительские) — количество_узлов(дочерние)
- сумма_расстояний(дочерние) = сумма_расстояний(родительские) + (всего_узлов — количество_узлов(дочерние)) — количество_узлов(дочерние)
- сумма_расстояний(дочерние) = сумма_расстояний(родительские) + общее_узел — 2*число_узлов(дочерние)
Следовательно, после того как мы вычислим сумму расстояний для одного корневого узла, используя обход в обратном порядке, мы можем пройти по дереву, используя обход в предварительном порядке, заполняя недостающие значения для некорневых узлов.
Это оставляет нам оптимальное решение O(n).
#include <unordered_map>
#include <vector>
struct TreeInfo {
TreeInfo(int n) : subtree_sum(n,0), node_count(n,0), result(n,0) {}
std::vector<int> subtree_sum;
std::vector<int> node_count;
std::vector<int> result;
};
void post_order(int node, int parent,
const std::unordered_multimap<int,int>& neighbours,
TreeInfo& info) {
// If there are no children we have zero distance and one node.
info.subtree_sum[node] = 0;
info.node_count[node] = 1;
auto [begin, end] = neighbours.equal_range(node);
for (auto [from, to] : std::ranges::subrange(begin, end)) {
// Avoid looping back to the node we came from.
if (to == parent) continue;
// post_order traversal, visit children first
post_order(to, node, neighbours, info);
// accumulate number of nodes and distances
info.subtree_sum[node] += info.subtree_sum[to] +
info.node_count[to];
info.node_count[node] += info.node_count[to];
}
}
void pre_order(int node, int parent,
const std::unordered_multimap<int,int>& neighbours,
TreeInfo& info) {
// For the root node the subtree_sum matches the result.
if (parent == -1) {
info.result[node] = info.subtree_sum[node];
} else {
// Otherwise, we can calculate the result from the parent,
// in pre-order we visit the parent before the children.
info.result[node] = info.result[parent] +
info.result.size() - 2*info.node_count[node];
}
// Now visit any children.
auto [begin, end] = neighbours.equal_range(node);
for (auto [from, to] : std::ranges::subrange(begin, end)) {
if (to == parent) continue;
pre_order(to, node, neighbours, info);
}
}
std::vector<int>
distances_in_tree(int n,
const std::unordered_multimap<int,int> neighbours) {
TreeInfo info(n);
// post-order pass to calculate subtree_sum and node_count
post_order(0,-1,neighbours,info);
// pre-order pass to calculate result
pre_order(0,-1,neighbours,info);
return info.result;
}
