ICPC 2020 - B. The Cost of Speed Limits

#include <bits/stdc++.h>
using namespace std;

namespace {

struct Result {
    int start = 0;
    vector<long long> dp;
    long long best = 0;
};

struct Frame {
    int v = 0;
    int next_child = 0;
    long long sign_cost = 0;
    long long sum_child_speeds = 0;
    int child_count = 0;
    bool has_sum = false;
    int no_sign_start = 0;
    vector<long long> sum_dp;
};

int n;
long long c;
vector<vector<pair<int, int>>> graph;
vector<vector<int>> children;
vector<int> parent_speed;
vector<int> parent_speed_index;
vector<int> degree_list;
vector<int> subtree_size;
vector<int> all_speeds;

void merge_child(Frame& frame, int child, Result child_result) {
    frame.sign_cost += child_result.best;
    frame.sum_child_speeds += parent_speed[child];

    if (!frame.has_sum) {
        frame.no_sign_start = child_result.start;
        frame.sum_dp = move(child_result.dp);
        frame.has_sum = true;
        return;
    }

    int new_start = max(frame.no_sign_start, child_result.start);
    int new_len = int(all_speeds.size()) - new_start;
    vector<long long> merged(new_len);

    int offset_sum = new_start - frame.no_sign_start;
    int offset_child = new_start - child_result.start;
    for (int i = 0; i < new_len; ++i) {
        merged[i] = frame.sum_dp[i + offset_sum] + child_result.dp[i + offset_child];
    }

    frame.sum_dp = move(merged);
    frame.no_sign_start = new_start;
}

Result finalize_node(const Frame& frame, long long& answer) {
    if (frame.v == 0 && frame.child_count == 0) {
        answer = 0;
        return {};
    }

    if (frame.child_count == 0) {
        int start = parent_speed_index[frame.v];
        return {start, vector<long long>(all_speeds.size() - start, 0), 0};
    }

    if (frame.v == 0) {
        answer = frame.sign_cost;
        for (int i = 0; i < int(frame.sum_dp.size()); ++i) {
            long long speed = all_speeds[frame.no_sign_start + i];
            long long no_sign_cost =
                frame.sum_dp[i] + 1LL * frame.child_count * speed - frame.sum_child_speeds;
            answer = min(answer, no_sign_cost);
        }
        return {};
    }

    int start = parent_speed_index[frame.v];
    vector<long long> dp(all_speeds.size() - start, frame.sign_cost);

    int begin = max(start, frame.no_sign_start);
    for (int idx = begin; idx < int(all_speeds.size()); ++idx) {
        long long speed = all_speeds[idx];
        long long no_sign_cost =
            frame.sum_dp[idx - frame.no_sign_start] + 1LL * frame.child_count * speed - frame.sum_child_speeds;
        dp[idx - start] = min(frame.sign_cost, no_sign_cost);
    }

    long long best = (1LL << 62);
    for (int idx = start; idx < int(all_speeds.size()); ++idx) {
        long long branch_cost = dp[idx - start] + (all_speeds[idx] - parent_speed[frame.v]);
        best = min(best, branch_cost);
    }

    return {start, move(dp), best};
}

long long run_dp() {
    long long answer = 0;
    vector<Frame> stack;
    stack.reserve(n);

    Frame root;
    root.v = 0;
    root.sign_cost = (degree_list[0] >= 2 ? 1LL * degree_list[0] * c : 0LL);
    root.child_count = int(children[0].size());
    stack.push_back(move(root));

    while (!stack.empty()) {
        Frame& current = stack.back();

        if (current.next_child < int(children[current.v].size())) {
            int child = children[current.v][current.next_child++];

            Frame child_frame;
            child_frame.v = child;
            child_frame.sign_cost = (degree_list[child] >= 2 ? 1LL * degree_list[child] * c : 0LL);
            child_frame.child_count = int(children[child].size());
            stack.push_back(move(child_frame));
            continue;
        }

        Result result = finalize_node(current, answer);
        int finished = current.v;
        stack.pop_back();

        if (stack.empty()) {
            break;
        }

        merge_child(stack.back(), finished, move(result));
    }

    return answer;
}

void build_rooted_tree() {
    vector<int> parent(n, -1);
    vector<int> order;
    order.reserve(n);
    order.push_back(0);
    parent[0] = 0;

    for (int i = 0; i < int(order.size()); ++i) {
        int v = order[i];
        for (const auto& edge : graph[v]) {
            int to = edge.first;
            int speed = edge.second;
            if (parent[to] != -1) {
                continue;
            }
            parent[to] = v;
            parent_speed[to] = speed;
            children[v].push_back(to);
            order.push_back(to);
        }
    }

    subtree_size.assign(n, 1);
    for (int i = n - 1; i >= 1; --i) {
        int v = order[i];
        subtree_size[parent[v]] += subtree_size[v];
    }

    for (int v = 0; v < n; ++v) {
        sort(children[v].begin(), children[v].end(), [&](int lhs, int rhs) {
            return subtree_size[lhs] > subtree_size[rhs];
        });
    }
}

void solve() {
    cin >> n >> c;

    if (n == 1) {
        cout << 0 << '\n';
        return;
    }

    graph.assign(n, {});
    children.assign(n, {});
    parent_speed.assign(n, 0);
    parent_speed_index.assign(n, 0);
    degree_list.assign(n, 0);
    all_speeds.clear();
    all_speeds.reserve(n - 1);

    for (int i = 0; i < n - 1; ++i) {
        int u, v, s;
        cin >> u >> v >> s;
        --u;
        --v;

        graph[u].push_back({v, s});
        graph[v].push_back({u, s});
        ++degree_list[u];
        ++degree_list[v];
        all_speeds.push_back(s);
    }

    sort(all_speeds.begin(), all_speeds.end());
    all_speeds.erase(unique(all_speeds.begin(), all_speeds.end()), all_speeds.end());

    for (int v = 1; v < n; ++v) {
        parent_speed_index[v] = 0;
    }

    build_rooted_tree();
    for (int v = 1; v < n; ++v) {
        parent_speed_index[v] =
            int(lower_bound(all_speeds.begin(), all_speeds.end(), parent_speed[v]) - all_speeds.begin());
    }

    cout << run_dp() << '\n';
}

}  // namespace

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    solve();
    return 0;
}

\documentclass[11pt]{article}
\usepackage[margin=1in]{geometry}
\usepackage[T1]{fontenc}
\usepackage[utf8]{inputenc}
\usepackage{amsmath,amssymb,amsthm}
\usepackage{enumitem}

\title{ICPC World Finals 2020\\B. The Cost of Speed Limits}
\author{}
\date{}

\begin{document}
\maketitle

\section*{Problem Summary}

We are given a tree whose edges already have speed limits.
We may:
\begin{itemize}[leftmargin=*]
    \item install signs at an intersection, paying $c$ dollars per incident road there if the final speed limits around
    that intersection are not all equal;
    \item increase any road speed from $s$ to $t \ge s$, paying $t-s$.
\end{itemize}

We must minimize the total upgrade cost plus the total sign cost.

\section*{Key Observations}

\begin{itemize}[leftmargin=*]
    \item If an edge is finally assigned some speed $x$, and $x$ is not one of the original edge speeds, we can lower it
    to the largest original speed that is still at most $x$ without breaking any equality requirement. This never
    increases the cost.
    Therefore we only need to consider the sorted list of distinct original speeds
    \[
        W_0 < W_1 < \dots < W_{m-1},
    \]
    where $m \le n-1$.
    \item Root the tree at an arbitrary vertex.
    For a non-root vertex $v$ with parent edge speed $s_v$, define
    \[
        dp_v[i]
    \]
    as the minimum cost inside the subtree of $v$, \emph{excluding} the upgrade cost of the parent edge, under the
    condition that the final speed on the parent edge equals $W_i$.
    This state is only meaningful when $W_i \ge s_v$.
    \item For every child $u$ of $v$, if we force the edge $(v,u)$ to end at speed $W_i$, then the contribution of that
    whole child branch is
    \[
        (W_i - s_{vu}) + dp_u[i].
    \]
    We also need
    \[
        best_u = \min_i \bigl((W_i - s_{vu}) + dp_u[i]\bigr),
    \]
    which is the cheapest possible total contribution of child $u$ when $v$ decides to place a sign.
    \item At a vertex $v$ there are exactly two cases:
    \begin{itemize}[leftmargin=*]
        \item \textbf{Place a sign at $v$.}
        Then all child branches are independent.
        \item \textbf{Do not place a sign at $v$.}
        Then every incident edge of $v$ must end with the same speed, so every child edge must also become $W_i$.
    \end{itemize}
\end{itemize}

\section*{Algorithm}

\begin{enumerate}[leftmargin=*]
    \item Compress all original edge speeds into the sorted array $W$.
    \item Root the tree and process it with DFS from the leaves upward.
    \item For a leaf $v$, $dp_v[i] = 0$ for every feasible $i$, because there is no cost inside its subtree once the
    parent edge speed is fixed.
    \item For an internal non-root vertex $v$, let its children be $u_1,\dots,u_k$.
    Compute
    \[
        sign(v) =
        \begin{cases}
            \deg(v)\,c + \sum best_{u_j}, & \deg(v) \ge 2, \\
            \sum best_{u_j}, & \deg(v) \le 1.
        \end{cases}
    \]
    Then for every feasible $i$:
    \[
        nosign_v(i) = \sum_{j=1}^{k} \bigl((W_i - s_{vu_j}) + dp_{u_j}[i]\bigr),
    \]
    and
    \[
        dp_v[i] = \min(sign(v), nosign_v(i)).
    \]
    Finally,
    \[
        best_v = \min_i \bigl((W_i - s_v) + dp_v[i]\bigr).
    \]
    \item The root has no parent edge, so its answer is simply
    \[
        \min\left(sign(root), \min_i \sum_j \bigl((W_i - s_{root,u_j}) + dp_{u_j}[i]\bigr)\right).
    \]
    \item In the implementation, every vertex stores only the tail of its table starting from the first feasible speed.
    When merging children, we only keep the common tail that is feasible for all of them.
\end{enumerate}

\section*{Correctness Proof}

We prove that the algorithm returns the correct answer.

\paragraph{Lemma 1.}
There exists an optimal solution in which every final edge speed is one of the original edge speeds from the input.

\paragraph{Proof.}
Take any optimal solution and any edge whose final speed is $x$.
If $x$ is already an original speed, do nothing.
Otherwise let $y$ be the largest original speed with $y < x$.
All constraints in the problem depend only on equalities between incident edge speeds and on the requirement that no
edge is decreased below its original speed.
If a set of edges was all equal to $x$, changing all of them to $y$ keeps them equal and still keeps every edge at
least its original speed, because every original speed in that set is at most $y$ by construction.
The total upgrade cost strictly decreases or stays the same.
Repeating this argument removes every non-original final speed. \qed

\paragraph{Lemma 2.}
For a non-root vertex $v$ and a feasible speed $W_i$, if we place a sign at $v$, then the minimum possible cost inside
the subtree of $v$ equals $sign(v)$.

\paragraph{Proof.}
If $v$ has a sign, then the speeds on different child edges are independent; they no longer need to match the parent
edge or each other.
For each child $u$, the cheapest possible contribution of that whole branch is exactly $best_u$ by definition.
The sign itself costs $\deg(v)\,c$ when $\deg(v)\ge 2$, and zero otherwise.
Summing these independent costs gives $sign(v)$, and no cheaper solution exists because every child branch is already
optimized independently. \qed

\paragraph{Lemma 3.}
For a non-root vertex $v$ and a feasible speed $W_i$, if we do not place a sign at $v$, then the minimum possible cost
inside the subtree of $v$ equals $nosign_v(i)$.

\paragraph{Proof.}
Without a sign at $v$, all incident final speeds at $v$ must be equal.
Since the parent edge has speed $W_i$, every child edge must also end at speed $W_i$.
For a child $u$, upgrading edge $(v,u)$ to $W_i$ costs $W_i - s_{vu}$, and the remaining optimal cost inside $u$'s
subtree is exactly $dp_u[i]$ by definition of the state.
Different child subtrees are disjoint, so their costs add.
Hence the total minimum cost is precisely $nosign_v(i)$. \qed

\paragraph{Lemma 4.}
For every non-root vertex $v$ and feasible speed $W_i$, the recurrence
\[
    dp_v[i] = \min(sign(v), nosign_v(i))
\]
is correct.

\paragraph{Proof.}
Every valid solution for the subtree of $v$ falls into exactly one of the two cases from Lemma 2 and Lemma 3:
either $v$ has a sign or it does not.
The two lemmas give the optimal cost for each case, so taking their minimum yields the true optimum for state
$dp_v[i]$. \qed

\paragraph{Lemma 5.}
For every non-root vertex $v$, the value
\[
    best_v = \min_i \bigl((W_i - s_v) + dp_v[i]\bigr)
\]
is exactly the minimum total contribution of the whole branch rooted at $v$ to its parent.

\paragraph{Proof.}
If the final speed on the parent edge is $W_i$, then its own upgrade cost is $W_i - s_v$, and the remainder of the
branch costs $dp_v[i]$.
Taking the minimum over all feasible $i$ gives exactly the cheapest possible total contribution of that branch. \qed

\paragraph{Theorem.}
The algorithm outputs the correct answer for every valid input.

\paragraph{Proof.}
By Lemma 1, restricting the DP to the compressed speed set loses no optimal solution.
For every non-root vertex, Lemma 4 computes each DP state correctly, and Lemma 5 computes the correct branch summary
$best_v$ that its parent needs.
Therefore the DFS computes correct values bottom-up for the whole rooted tree.
At the root, the same two cases remain: either place a sign and pay $sign(root)$, or require all incident edges to be
equal to one compressed speed $W_i$ and pay the corresponding no-sign sum.
The algorithm returns the minimum of exactly these possibilities, so the final answer is optimal. \qed

\section*{Complexity Analysis}

Let $m$ be the number of distinct original speeds, so $m \le n-1$.

Each vertex processes at most $m$ DP states, and each child merge is linear in the length of the common feasible tail.
Thus the running time is $O(nm)$, which is $O(n^2)$ in the worst case.

The implementation stores only tail tables and merges children immediately, so it does not need to keep a full
$n \times m$ table in memory at once.

\section*{Implementation Notes}

\begin{itemize}[leftmargin=*]
    \item The DP state excludes the upgrade cost of the parent edge. This keeps the recurrence clean and makes
    $best_v$ the only summary value needed by the parent.
    \item States with $W_i$ below the original parent-edge speed are impossible and are not stored.
    \item Sorting children by decreasing subtree size keeps the DFS memory usage low in long skinny trees.
    \item Use \texttt{long long}: the answer can exceed $2^{31}-1$.
\end{itemize}

\end{document}