Game on Axis (Codeforces 1787D)
Tags: Implement, dsu

There are $$$n$$$ points $$$1,2,\ldots,n$$$, each point $$$i$$$ has a number $$$a_i$$$ on it. You're playing a game on them. Initially, you are at point $$$1$$$. When you are at point $$$i$$$, take following steps:  If $$$1\le i\le n$$$, go to $$$i+a_i$$$,  Otherwise, the game ends. Before the game begins, you can choose two integers $$$x$$$ and $$$y$$$ satisfying $$$1\le x\le n$$$, $$$-n \le y \le n$$$ and replace $$$a_x$$$ with $$$y$$$ (set $$$a_x := y$$$). Find the number of distinct pairs $$$(x,y)$$$ such that the game that you start after making the change ends in a finite number of steps.Notice that you do not have to satisfy $$$a_x\not=y$$$.

Input: Each test contains multiple test cases. The first line contains an integer $$$t$$$ ($$$1\le t\le 10^4)$$$ — the number of test cases.The first line of each test case contains one integer $$$n$$$ ($$$1\le n\le 2\cdot 10^5$$$) — the number of points.The second line contains $$$n$$$ integers $$$a_1,a_2,\ldots,a_n$$$ ($$$-n \le a_i \le n$$$) — the numbers on the axis.It's guaranteed that the sum of $$$n$$$ does not exceed $$$2\cdot 10^5$$$.

Output: For each test case, print a line containing a single integer — the number of distinct pairs $$$(x,y)$$$ with which the game ends.

Note: In the first test case, the pairs $$$(x,y)$$$ with which the game ends are $$$(1,-1)$$$ and $$$(1,1)$$$, corresponding to the routes $$$1\rightarrow 0$$$ and $$$1\rightarrow 2$$$. Note that $$$(1,2)$$$ is invalid since when $$$n=1$$$, $$$y=2$$$ violates $$$-n\le y\le n$$$. $$$(1,0)$$$ is also invalid since you will go from $$$1$$$ to $$$1$$$ forever.In the second test case, the pairs are $$$(1,-2),(1,-1),(1,2),(2,-2),(2,-1),(2,0),(2,1),(2,2)$$$.In the fourth test case, the pairs are $$$(1,-2),(1,-1),(1,1),(1,2),(2,-2),(2,1),(2,2)$$$.

Sample Inputs:
9
1
0
2
-1 0
2
1 -1
2
1 1
3
-1 -2 -1
3
1 -2 -1
4
-1 4 -2 1
5
1 1 1 1 -4
5
1 1 1 1 1

Sample Outputs:
2
8
6
7
20
17
34
30
40

Solution (Python):
def solve():
    n = int(input())
    nums = list(map(int, input().split()))
    ans = n * (n * 2 + 1)
    visited = [0] * n
    path = []
    pt = 0
    flag = False
    while not visited[pt]:
        path.append(pt)
        visited[pt] = 1
        pt += nums[pt]
        if not 0 <= pt < n:
            break
    else:
        ans -= (n - sum(visited)) * (2 * n + 1)
        flag = True
    cnt = sum(visited)
    if flag:
        ans -= cnt * cnt
    else:
        ans -= (1 + cnt) * cnt // 2
    new_visited = [0] * n
    note = visited[:]
    for i in range(n):
        if not note[i]:
            tmp = []
            f = False
            while not note[i]:
                tmp.append(i)
                note[i] = 1
                i += nums[i]
                if not 0 <= i < n:
                    break
                if note[i] and new_visited[i] >= 0 or visited[i]:
                    f = True
            if f:
                if new_visited[i]:
                    i = new_visited[i] - 1
                for x in tmp:
                    new_visited[x] = i + 1
            else:
                for x in tmp:
                    new_visited[x] = -1
    path_note = {v: i for (i, v) in enumerate(path)}
    for i in range(n):
        if not visited[i]:
            if new_visited[i] != -1:
                if visited[new_visited[i] - 1] == 0 or flag:
                    ans -= cnt
                else:
                    ans -= cnt - path_note[new_visited[i] - 1]
    print(ans)
for _ in range(int(input())):
    solve()

Solution (Python):
def solve():
    n = int(input())
    nums = list(map(int, input().split()))
    ans = n * (n * 2 + 1)
    visited = [0] * n
    path = []
    pt = 0
    flag = False
    while not visited[pt]:
        path.append(pt)
        visited[pt] = 1
        pt += nums[pt]
        if not 0 <= pt < n:
            break
    else:
        ans -= (n - sum(visited)) * (2 * n + 1)
        flag = True
    cnt = sum(visited)
    if flag:
        ans -= cnt * cnt
    else:
        ans -= (1 + cnt) * cnt // 2
    new_visited = [0] * n
    note = visited[:]
    for i in range(n):
        if not note[i]:
            tmp = []
            f = False
            while not note[i]:
                tmp.append(i)
                note[i] = 1
                i += nums[i]
                if not 0 <= i < n:
                    break
                if note[i] and new_visited[i] >= 0 or visited[i]:
                    f = True
            if f:
                if new_visited[i]:
                    i = new_visited[i] - 1
                for x in tmp:
                    new_visited[x] = i + 1
            else:
                for x in tmp:
                    new_visited[x] = -1
    path_note = {v: i for (i, v) in enumerate(path)}
    for i in range(n):
        if not visited[i]:
            if new_visited[i] != -1:
                if visited[new_visited[i] - 1] == 0 or flag:
                    ans -= cnt
                else:
                    ans -= cnt - path_note[new_visited[i] - 1]
    print(ans)
for _ in range(int(input())):
    solve()

Solution (Python):
import math
import random
import time
from collections import defaultdict
contor = 0
matricea = {}
alfabet = {'a': 1, 'b': 2, 'c': 3, 'd': 4, 'e': 5, 'f': 6, 'g': 7, 'h': 8, 'i': 9, 'j': 10, 'k': 11, 'l': 12, 'm': 13, 'n': 14, 'o': 15, 'p': 16, 'q': 17, 'r': 18, 's': 19, 't': 20, 'u': 21, 'v': 22, 'w': 23, 'x': 24, 'y': 25, 'z': 26}
alfabet_2 = {'1': 'a', '2': 'b', '3': 'c', '4': 'd', '5': 'e', '6': 'f', '7': 'g', '8': 'h', '9': 'i', '10': 'j', '11': 'k', '12': 'l', '13': 'm', '14': 'n', '15': 'o', '16': 'p', '17': 'q', '18': 'r', '19': 's', '20': 't', '21': 'u', '22': 'v', '23': 'w', '24': 'x', '25': 'y', '26': 'z'}
from bisect import bisect_left as lower_bound
from bisect import bisect_right as upper_bound

class FenwickTree:

    def __init__(self, x):
        bit = self.bit = list(x)
        size = self.size = len(bit)
        for i in range(size):
            j = i | i + 1
            if j < size:
                bit[j] += bit[i]

    def update(self, idx, x):
        """"""
        while idx < self.size:
            self.bit[idx] += x
            idx |= idx + 1

    def __call__(self, end):
        """"""
        x = 0
        while end:
            x += self.bit[end - 1]
            end &= end - 1
        return x

    def find_kth(self, k):
        """"""
        idx = -1
        for d in reversed(range(self.size.bit_length())):
            right_idx = idx + (1 << d)
            if right_idx < self.size and self.bit[right_idx] <= k:
                idx = right_idx
                k -= self.bit[idx]
        return (idx + 1, k)

class SortedList:
    block_size = 700

    def __init__(self, iterable=()):
        self.macro = []
        self.micros = [[]]
        self.micro_size = [0]
        self.fenwick = FenwickTree([0])
        self.size = 0
        for item in iterable:
            self.insert(item)

    def insert(self, x):
        i = lower_bound(self.macro, x)
        j = upper_bound(self.micros[i], x)
        self.micros[i].insert(j, x)
        self.size += 1
        self.micro_size[i] += 1
        self.fenwick.update(i, 1)
        if len(self.micros[i]) >= self.block_size:
            self.micros[i:i + 1] = (self.micros[i][:self.block_size >> 1], self.micros[i][self.block_size >> 1:])
            self.micro_size[i:i + 1] = (self.block_size >> 1, self.block_size >> 1)
            self.fenwick = FenwickTree(self.micro_size)
            self.macro.insert(i, self.micros[i + 1][0])

    def pop(self, k=-1):
        (i, j) = self._find_kth(k)
        self.size -= 1
        self.micro_size[i] -= 1
        self.fenwick.update(i, -1)
        return self.micros[i].pop(j)

    def __getitem__(self, k):
        (i, j) = self._find_kth(k)
        return self.micros[i][j]

    def count(self, x):
        return self.upper_bound(x) - self.lower_bound(x)

    def __contains__(self, x):
        return self.count(x) > 0

    def lower_bound(self, x):
        i = lower_bound(self.macro, x)
        return self.fenwick(i) + lower_bound(self.micros[i], x)

    def upper_bound(self, x):
        i = upper_bound(self.macro, x)
        return self.fenwick(i) + upper_bound(self.micros[i], x)

    def _find_kth(self, k):
        return self.fenwick.find_kth(k + self.size if k < 0 else k)

    def __len__(self):
        return self.size

    def __iter__(self):
        return (x for micro in self.micros for x in micro)

    def __repr__(self):
        return str(list(self))

def prime_generator(nr_elemente_prime):
    vector_prime = [-1] * nr_elemente_prime
    vector_rasp = [0] * nr_elemente_prime
    vector_prime[1] = 1
    vector_rasp[1] = 1
    contor = 2
    for i in range(2, nr_elemente_prime):
        if vector_prime[i] == -1:
            vector_prime[i] = 1
            vector_rasp[contor] = i
            contor = contor + 1
        for j in range(i + i, nr_elemente_prime, i):
            if vector_prime[j] == -1:
                vector_prime[j] = i
    set_prime = set(vector_rasp)
    set_prime.remove(0)
    set_prime.remove(1)
    answ_prime = list(set_prime)
    answ_prime.sort()
    return answ_prime

def binary_search(vectorul, valoarea):
    left = 0
    right = len(vectorul) - 1
    while left <= right:
        centru = (left + right) // 2
        if vectorul[centru] < valoarea:
            left = centru + 1
        else:
            right = centru - 1
    return len(vectorul) - left

def functie_combinari(vectorul, pozitie, sfarsit):
    global contor
    global matricea
    if pozitie >= sfarsit:
        setul = set()
        for hh in vectorul:
            setul.add(hh)
        if len(setul) == sfarsit:
            contor += 1
            matricea[contor] = vectorul
            print('contor=', contor)
    else:
        for vectorul[pozitie] in range(sfarsit):
            functie_combinari(vectorul, pozitie + 1, sfarsit)

def main():
    vec_answ = []
    pp = 10 ** 40
    vec_answ = []
    lista_prime = prime_generator(4 * 10 ** 4)
    teste = int(input())
    for gg in range(teste):
        n = int(input())
        lista = list(map(int, input().split()))
        set_dr_principal = set()
        set_dr_iesire = set()
        set_dr_blocate = set()
        set_dau_in_principal = set()
        set_dr_principal.add(0)
        sume = [0] * len(set_dr_principal)
        dictionar_principal = {}
        game = 0
        pozitie = 0
        next = lista[pozitie] + pozitie
        dictionar_principal[0] = 0
        contor = 1
        while next >= 0 and next < n and (next not in set_dr_principal):
            dictionar_principal[next] = contor
            contor += 1
            set_dr_principal.add(next)
            next = lista[next] + next
        if next >= n or next < 0:
            game = 1
        dictionar = {}
        sume = [1] * len(set_dr_principal)
        pozitie = 1
        while pozitie < n:
            set_actual = set()
            initial = pozitie
            set_actual.add(initial)
            posibil = 1
            if pozitie in set_dr_principal or pozitie in set_dr_iesire or pozitie in set_dr_blocate or (pozitie in set_dau_in_principal):
                posibil = 0
            while posibil == 1:
                if initial < 0 or initial >= n:
                    posibil = 0
                    for t in set_actual:
                        set_dr_iesire.add(t)
                elif initial in set_dr_principal:
                    posibil = 0
                    sume[dictionar_principal[initial]] += len(set_actual)
                    for t in set_actual:
                        dictionar[t] = initial
                        set_dau_in_principal.add(t)
                elif initial in set_dau_in_principal:
                    posibil = 0
                    unde = dictionar[initial]
                    val = dictionar_principal[unde]
                    sume[val] += len(set_actual)
                    for t in set_actual:
                        dictionar[t] = unde
                        set_dau_in_principal.add(t)
                elif initial in set_dr_iesire:
                    posibil = 0
                    for t in set_actual:
                        set_dr_iesire.add(t)
                elif initial in set_dr_blocate or (initial != pozitie and initial in set_actual) or lista[pozitie] == 0:
                    for t in set_actual:
                        set_dr_blocate.add(t)
                    posibil = 0
                else:
                    set_actual.add(initial)
                    initial += lista[initial]
            pozitie += 1
        answ = 0
        sume_generale = [0] * len(set_dr_principal)
        sume_generale[0] = sume[0]
        for i in range(1, len(set_dr_principal)):
            sume_generale[i] = sume[i] + sume_generale[i - 1]
        suma_totala = sume_generale[len(set_dr_principal) - 1]
        if game == 0:
            answ = len(set_dr_principal) * (n + 1 + len(set_dr_iesire))
        else:
            for i in range(len(set_dr_principal)):
                answ += n + 1 + suma_totala - sume_generale[i] + len(set_dr_iesire)
            answ += (2 * n + 1) * (len(set_dr_blocate) + len(set_dr_iesire) + len(set_dau_in_principal))
        vec_answ.append(answ)
    print(*vec_answ, sep='\n')
main()