AtCoder - 167D. Teleporter

You can practice the problem here.

Problem statement

The Kingdom of Takahashi has N towns, numbered 1 through N.

There is one teleporter in each town. The teleporter in Town i (1 <= i <= N) sends you to Town A[i].

Takahashi, the king, loves the positive integer K. The selfish king wonders what town he will be in if he starts at Town 1 and uses a teleporter exactly K times from there.

Help the king by writing a program that answers this question.

Sample Input 1

4 5
3 2 4 1

Sample Output 1

4

Sample Input 2

6 727202214173249351
6 5 2 5 3 2

Sample Output 2

2

Solution 1: Math

If you take a look at the sample input 2, we can see that there is no way to traverse k times by listing out the travel. We can easily obverse the pattern the loop starts at a certain point.

Example:

6 5 2 5 3 2
1 -> 6 -> 2 -> 5 -> 3 -> 2 -> 5 -> 3 -> 2 -> 5 ....  

Starting from the third step, it starts the loop 2 -> 5 -> 3.

We can first create a map to store the index (0-based) and the value.

m[i] = a[i]

Then we can find how many steps we need to make before entering the loop, let's say c[i], and how many steps to reach the first duplicate element, let's say step.

bool duplicate = false;
ll i = 0, step = 0, kk = k;
while(!duplicate && kk){
    step++; // step is used to find out when a town has been visited
    if(d[i] == 1){
        // if it is visited before, that is the start of the loop
        duplicate = true;
        continue;
    }
    d[i] = 1; // set the current i to 1, in other word, this has been visited
    c[i] = step; // we need to store the step to another map because we need to calculate from where the looping starts
    i = m[i] - 1; // -1 because it s 0 based
    kk--; // monitor if it s out of boundary
}

We can calculate how many times we need to teleport from the repeating pattern as we know the length of the repeating pattern till k k - c[i], and the length of the repeating pattern step - c[i].

6 5 2 5 3 2
1 -> 6 -> 2 -> 5 -> 3 -> 2 -> 5 -> 3 -> 2 -> 5 ....
          x              x 
|-----------step---------|
          |--------------k-c[i]---------------- ....
|---c[i]--|                         
|------------------------k--------------------- ....

It performs only if it is not out of boundary.

if(kk){
    ll r = (k - c[i]) % (step - c[i]);
    while(r >= 0){
        i = m[i] - 1;
        r--;
    } 
}
OUT(i + 1);

Solution 2: Binary Lifting

After half a year, I came up with another solution using Binary Lifting which is a technique used to find the k-th ancestor of any node in a tree in O(logn). You can use it to find out Lowest Common Ancestor(LCA) as well.

We know the k can go up to 10 ^ 18, which means we only need around 60 nodes (log2(10 ^ 18)).

const int mxN = 200005;
const int mxNode = 60; // ~ log2(k) 
int up[mxN][mxNode];

First we need to preprecess the tree in O(nlogn) using dynamic programming.

void binary_lifting() {
  for(int i = 1; i < mxNode; i++) {
    for(int u = 0; u < mxN; u++) {
      up[u][i] = up[up[u][i - 1]][i - 1];
    }
  }
}

Then we can easily find the k-th ancestor of any node in a tree in O(logn). In this case, we start from node 0.

int kth_ancestor(ll u, ll k) {
  for(int i = 0; i < mxNode; i++) {
    if(k & (1LL << i)) u = up[u][i];
  }
  return u;
}

The answer would be kth_ancestor(0, k) + 1.

Source Code

The full solution is avilable here.