Problem Statement:

Codeforces Round #353 (Div. 2) - E. Trains and StatisticSummary:

A graph with node 1..n such that node i have an edge to every nodes in i+1 to a[i] inclusive (where i+1 <= a[i] <= n). Calculate the sum of all the length shortest path from every pair of nodes.

Solution:

Another very interesting problem. The solution boils down to the following observation: to reach a destination with the least tickets, from each station we where we start/alight, buy a ticket to the station that can reach the furthest station. (I think this is actually non-trivial, although now it feels kind of intuitive. But yeah, I didn't come up with this by myself).

Mathematically speaking, consider station i. Say j is in i+1 to a[i] such that a[j] is the largest amongst the rest of the station in that range (if there are several such j, pick any one of them). The claim is, the path i->j->k for every k > a[i] is a shortest path from i to k.

Proof by contradiction: suppose otherwise, then there exists another station x != j directly reachable form i such that i->x->k has a shorter path than i->j->k for some k, but a[x] < a[j]. If k is directly reachable from x (i.e. k <= a[x]), then since a[x] < a[j], k is also directly reachable from j, so x->k and j->k have the same length = 1, and hence i->x->k and i->j->k are of the same length too (hence for this case, i->j->k is as good as i->x->k). Hence k > a[x]. We conclude k must be larger than a[j] too because otherwise i->j->k will need only 2 tickets while i->x->k needs more than 2. So for k > a[j], we know that there exists another station directly reachable from x, say y (i.e. y <= a[x]) such that i->x->y->k is shorter than i->j->k. However, if such y exists, y is also directly reachable from j, so that means we also have i->j->y->k which then concludes that i->j->k is at least as good as i->x->k, the shortest path.

So now what we need to do is to compute the sum of length of shortest paths for each individual station from right to left. Let sum[i] be the sum of length of all shortest paths from i to i+1, i+2, ..., n. Say we have computed sum[i+1], sum[i+2], ..., sum[n-1]. To compute sum[i], we break it down into two parts:

1. For stations i+1 to a[i]:

Each is directly reachable, so the sum is a[i]-i.

2. For all stations k > a[i]:

We first find j the station directly reachable from i that has largest a[j]. Then for all k > a[i], the length of shortest path from i to k is 1 + the length of shortest path from j to k. sum[j] is the sum of all shortest path from j+1 to n. If we remove the length counts from j+1 to a[i] from sum[j] (which is a[i]-j), the sum of lengths for k > a[i] is sum[j] - (a[i] - j) + (n - a[i]).

Hence sum[i] = sum[j] + j - i + n - a[i]. (By the way, notice that the computation was done using 1-index instead of 0-index).

To find the largest j for each i, we can use a segment tree. So the complexity of the solution is O(n lg n). This problem is good. And now I know how to save more by spending less on commute.

#include <iostream> #include <cstdio> #include <vector> #include <algorithm> #include <utility> using namespace std; int n; vector<int> a, tree; vector<int64_t> dp; void Build(int L=1, int R=n-1, int p=1) { if (L>R) return; if (L==R) { tree[p] = L; return; } int M=(L+R)/2; Build(L,M,p*2); Build(M+1,R,2*p+1); tree[p] = a[tree[2*p]] > a[tree[2*p+1]] ? tree[2*p] : tree[2*p+1]; } int Rmq(int S, int T, int L=1,int R=n-1, int p=1) { if (T<L||R<S) return -1; if (S<=L&&R<=T) return tree[p]; int M = (L+R)/2; int left = Rmq(S, T, L, M, 2*p); int right = Rmq(S, T, M+1, R, 2*p+1); if (left == -1) return right; if (right == -1) return left; return a[left] > a[right] ? left : right; } int main() { scanf("%d",&n); a.resize(n); dp.resize(n); for(int i=1;i<n;++i){ scanf("%d",&a[i]); a[i]; } dp[n-1] = 1; tree.resize(4*n); Build(); int64_t ans = 1; for(int i=n-2;i>=1;--i){ int j = Rmq(i+1, a[i]); dp[i] = dp[j] + n-i - a[i] + j; ans += dp[i]; } cout << ans << endl; return 0; }