Problem: Maximize Value of Function in a Ball Passing Game

You are given an integer array `receiver` of length `n` and an integer `k`.
`n` players are playing a ball-passing game.

You choose the starting player, `i`. The game proceeds as follows: player `i`
passes the ball to player `receiver[i]`, who then passes it to
`receiver[receiver[i]]`, and so on, for `k` passes in total. The game's score
is the sum of the indices of the players who touched the ball, including
repetitions, i.e. `i + receiver[i] + receiver[receiver[i]] + ... +
receiver(k)[i]`.

Return the maximum possible score.

Notes:

* `receiver` may contain duplicates.
* `receiver[i]` may be equal to `i`.


Python solution:
import math
import itertools
import bisect
import collections
import string
import heapq
import functools
import sortedcontainers
from typing import List, Dict, Tuple, Iterator

class Solution:
  def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:
    n = len(receiver)
    m = int(math.log2(k)) + 1
    ans = 0
    jump = [[0] * m for _ in range(n)]
    summ = [[0] * m for _ in range(n)]

    for i in range(n):
      jump[i][0] = receiver[i]
      summ[i][0] = receiver[i]

    for j in range(1, m):
      for i in range(n):
        midNode = jump[i][j - 1]
        jump[i][j] = jump[midNode][j - 1]
        summ[i][j] = summ[i][j - 1] + summ[midNode][j - 1]

    for i in range(n):
      currSum = i
      currPos = i
      for j in range(m):
        if (k >> j) & 1 == 1:
          currSum += summ[currPos][j]
          currPos = jump[currPos][j]
      ans = max(ans, currSum)

    return ans


Java solution:
class Solution {
  public long getMaxFunctionValue(List<Integer> receiver, long k) {
    final int n = receiver.size();
    final int m = (int) (Math.log(k) / Math.log(2)) + 1;
    long ans = 0;
    // jump[i][j] := the the node you reach after jumping 2^j steps from i
    int[][] jump = new int[n][m];
    // sum[i][j] := the sum of the first 2^j nodes you reach when jumping from i
    long[][] sum = new long[n][m];

    for (int i = 0; i < n; ++i) {
      jump[i][0] = receiver.get(i);
      sum[i][0] = receiver.get(i);
    }

    // Calculate binary lifting.
    for (int j = 1; j < m; ++j)
      for (int i = 0; i < n; ++i) {
        final int midNode = jump[i][j - 1];
        //   the the node you reach after jumping 2^j steps from i
        // = the node you reach after jumping 2^(j - 1) steps from i
        // + the node you reach after jumping another 2^(j - 1) steps
        jump[i][j] = jump[midNode][j - 1];
        //   the sum of the first 2^j nodes you reach when jumping from i
        // = the sum of the first 2^(j - 1) nodes you reach when jumping from i
        // + the sum of another 2^(j - 1) nodes you reach
        sum[i][j] = sum[i][j - 1] + sum[midNode][j - 1];
      }

    for (int i = 0; i < n; ++i) {
      long currSum = i;
      int currPos = i;
      for (int j = 0; j < m; ++j)
        if ((k >> j & 1) == 1) {
          currSum += sum[currPos][j];
          currPos = jump[currPos][j];
        }
      ans = Math.max(ans, currSum);
    }

    return ans;
  }
}


C++ solution:
class Solution {
 public:
  long long getMaxFunctionValue(vector<int>& receiver, long long k) {
    const int n = receiver.size();
    const int m = log2(k) + 1;
    long ans = 0;
    // jump[i][j] := the the node you reach after jumping 2^j steps from i
    vector<vector<int>> jump(n, vector<int>(m));
    // sum[i][j] := the sum of the first 2^j nodes you reach when jumping from i
    vector<vector<long>> sum(n, vector<long>(m));

    for (int i = 0; i < n; ++i) {
      jump[i][0] = receiver[i];
      sum[i][0] = receiver[i];
    }

    // Calculate binary lifting.
    for (int j = 1; j < m; ++j)
      for (int i = 0; i < n; ++i) {
        const int midNode = jump[i][j - 1];
        //   the the node you reach after jumping 2^j steps from i
        // = the node you reach after jumping 2^(j - 1) steps from i
        // + the node you reach after jumping another 2^(j - 1) steps
        jump[i][j] = jump[midNode][j - 1];
        //   the sum of the first 2^j nodes you reach when jumping from i
        // = the sum of the first 2^(j - 1) nodes you reach when jumping from i
        // + the sum of another 2^(j - 1) nodes you reach
        sum[i][j] = sum[i][j - 1] + sum[midNode][j - 1];
      }

    for (int i = 0; i < n; ++i) {
      long currSum = i;
      int currPos = i;
      for (int j = 0; j < m; ++j)
        if (k >> j & 1) {
          currSum += sum[currPos][j];
          currPos = jump[currPos][j];
        }
      ans = max(ans, currSum);
    }

    return ans;
  }
};