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

Problem Statement:

Codeforces Round #353 (Div. 2) - E. Trains and Statistic
Summary:
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;
}