Madoka and the Sixth-graders (Codeforces 1647E)
Tags: Greedy, dfs and similar

After the most stunning success with the fifth-graders, Madoka has been trusted with teaching the sixth-graders.There's $$$n$$$ single-place desks in her classroom. At the very beginning Madoka decided that the student number $$$b_i$$$ ($$$1 \le b_i \le n$$$) will sit at the desk number $$$i$$$. Also there's an infinite line of students with numbers $$$n + 1, n + 2, n + 3, \ldots$$$ waiting at the door with the hope of being able to learn something from the Madoka herself. Pay attention that each student has his unique number.After each lesson, the following happens in sequence.   The student sitting at the desk $$$i$$$ moves to the desk $$$p_i$$$. All students move simultaneously.  If there is more than one student at a desk, the student with the lowest number keeps the place, and the others are removed from the class forever.  For all empty desks in ascending order, the student from the lowest number from the outside line occupies the desk. Note that in the end there is exactly one student at each desk again. It is guaranteed that the numbers $$$p$$$ are such that at least one student is removed after each lesson. Check out the explanation to the first example for a better understanding.After several (possibly, zero) lessons the desk $$$i$$$ is occupied by student $$$a_i$$$. Given the values $$$a_1, a_2, \ldots, a_n$$$ and $$$p_1, p_2, \ldots, p_n$$$, find the lexicographically smallest suitable initial seating permutation $$$b_1, b_2, \ldots, b_n$$$.The permutation is an array of $$$n$$$ different integers from $$$1$$$ up to $$$n$$$ in any order. For example, $$$[2,3,1,5,4]$$$ is a permutation, but $$$[1,2,2]$$$ is not ($$$2$$$ occurs twice). $$$[1,3,4]$$$ is not a permutation either ($$$n=3$$$ but there's $$$4$$$ in the array).For two different permutations $$$a$$$ and $$$b$$$ of the same length, $$$a$$$ is lexicographically less than $$$b$$$ if in the first position where $$$a$$$ and $$$b$$$ differ, the permutation $$$a$$$ has a smaller element than the corresponding element in $$$b$$$.

Input: The first line of input data contains an integer $$$n$$$ ($$$2 \le n \le 10^5$$$) — a number of desks in the classroom.The second line contains $$$n$$$ integers $$$p_1, p_2, \ldots, p_n$$$ ($$$1 \leq p_i \leq n$$$) — desks where the students move. It is guaranteed that $$$p$$$ has at least two equal elements.The third line contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$1 \leq a_i \leq 10^9$$$) — the final seating of the students. It is guaranteed that there is an initial permutation from which the seating $$$a$$$ can be obtained.

Output: In the only line print $$$n$$$ integers $$$b_1, b_2, \ldots, b_n$$$ ($$$1 \le b_i \le n$$$) — lexicographically minimum permutation describing the initial seating of the sixth-graders that can lead to the final seating $$$a$$$.

Note: The description of the first test is below:  The first picture shows the starting permutation, which is the answer. Then the students sitting at desks $$$1, 2$$$ are transferred to a $$$5$$$ desk. Also, a $$$1$$$ student moved from a $$$5$$$ desk, and a student from a $$$4$$$ disk is transferred to a $$$3$$$ desk.Thus, after all these transfers permutation shown in the second image is obtained. Then, at the desk with the number $$$5$$$, the student with the number $$$3$$$ is expelled, and at the desk with the number $$$3$$$, the student with the number $$$5$$$ is expelled. (Since their numbers are not the smallest) Then new students with numbers $$$6, 7$$$ sit at desks numbered $$$2, 4$$$. And this permutation (after the end of the first lesson) is shown in the third image.The $$$4$$$ image shows the seating arrangement, after the second lesson before all the extra ones were kicked out. And the fifth shows the final seating after $$$2$$$ lesson.

Sample Inputs:
5
5 5 3 3 1
1 8 2 9 4

5
1 3 2 5 2
3 2 5 4 1

10
10 8 5 3 7 8 6 6 1 1
5 26 24 27 21 4 18 2 28 1

Sample Outputs:
1 3 2 5 4

3 2 5 4 1

5 4 2 6 7 8 3 9 1 10

Solution (Python):
import sys
input = sys.stdin.readline
n = int(input())
p = list(map(int, input().split()))
a = list(map(int, input().split()))

class Doubling:

    def __init__(self, k, prod, prod1):
        self.k = k
        self.prod = prod
        self.prod1 = prod1
        self.f = [prod1()]
        for i in range(k):
            self.f.append(prod(self.f[-1], self.f[-1]))

    def double(self, n):
        n -= 1
        res = self.prod1()
        i = 0
        while n:
            if n & 1:
                res = self.prod(res, self.f[i])
            n >>= 1
            i += 1
        return res

def prod(x, y):
    res = [0] * n
    for i in range(n):
        res[i] = y[x[i] - 1]
    return res

def prod1():
    return p
dbl = Doubling(31, prod, prod1)
k = max(a) - n
if k == 0:
    print(*a)
    exit(0)
k //= n - len(set(p))
pk = dbl.double(k)
used = [0] * (n + 1)
res = [0] * n
s = []
for i in range(n):
    j = pk[i] - 1
    if a[j] <= n and used[a[j]] == 0:
        j = pk[i] - 1
        res[i] = a[j]
        used[a[j]] = 1
ls = [0] * n
t = [[] for _ in range(n + 1)]
for i in range(n):
    if res[i]:
        continue
    j = pk[i] - 1
    ls[i] = a[j] + 1
    if ls[i] > n:
        t[0].append(i)
    else:
        t[ls[i]].append(i)
import heapq
hpush = heapq.heappush
hpop = heapq.heappop
xs = set(range(1, n + 1)) - set(res)
h = []
for x in range(n + 1):
    for i in t[x]:
        hpush(h, i)
    if x in xs:
        i = hpop(h)
        res[i] = x
print(*res)

Solution (Python):
import sys
I = lambda : [*map(int, sys.stdin.readline().split())]

class unionfind:

    def __init__(self, n):
        self.size = [1] * n
        self.parent = [None] * n
        self.rep = [i for i in range(n)]

    def find(self, x):
        path = []
        curr = x
        while self.parent[curr] is not None:
            path.append(curr)
            curr = self.parent[curr]
        for guy in path:
            self.parent[guy] = curr
        return curr

    def union(self, x, y):
        currx = self.find(x)
        curry = self.find(y)
        if currx == curry:
            return None
        if self.size[currx] < self.size[curry]:
            self.parent[currx] = curry
            self.size[curry] += self.size[currx]
            self.rep[curry] = max(self.rep[currx], self.rep[curry])
        else:
            self.parent[curry] = currx
            self.size[currx] += self.size[curry]
            self.rep[currx] = max(self.rep[currx], self.rep[curry])
(n,) = I()
p = I()
for i in range(n):
    p[i] -= 1
a = I()
for i in range(n):
    a[i] -= 1
m = max(a)
l = n - len(set(p))
rounds = (m - n + 1) // l
perms = [p]
for i in range(len(bin(rounds)) - 3):
    prev = perms[-1]
    new = [prev[prev[i]] for i in range(n)]
    perms.append(new)
curr = [i for i in range(n)]
for i in range(len(perms)):
    if rounds >> i & 1:
        new = [perms[i][curr[j]] for j in range(n)]
        curr = new
smls = {}
firsts = {}
for i in range(n):
    if curr[i] not in smls:
        smls[curr[i]] = a[curr[i]]
        firsts[curr[i]] = i
dsf = unionfind(n)
perm = [-1] * n
for guy in firsts:
    perm[firsts[guy]] = smls[guy]
    if smls[guy] + 1 < n:
        dsf.union(smls[guy], smls[guy] + 1)
for i in range(n):
    if perm[i] == -1:
        x = smls[curr[i]]
        y = dsf.find(x)
        use = dsf.rep[y]
        perm[i] = use
        if use + 1 < n:
            dsf.union(y, use + 1)
perm = [guy + 1 for guy in perm]
print(*perm)

Solution (Python):
import sys
input = sys.stdin.buffer.readline

class set_with_geq:

    def __init__(self, n):
        self.n = n + 2
        self.par = [i for i in range(self.n)]
        self.size = [1] * self.n
        self.val = [i for i in range(self.n)]

    def find_parent(self, a):
        if self.par[a] == a:
            return a
        self.par[a] = self.find_parent(self.par[a])
        return self.par[a]

    def geq(self, a):
        return self.val[self.find_parent(a)]

    def delete(self, a):
        b = self.find_parent(a + 1)
        a = self.find_parent(a)
        val = self.val[b]
        if a != b:
            if self.size[a] < self.size[b]:
                (a, b) = (b, a)
            self.par[b] = a
            self.size[a] += self.size[b]
            self.val[a] = val
n = int(input())
move = [[0] * n for i in range(30)]
move[0] = list(map(lambda x: int(x) - 1, input().split()))
for j in range(1, 30):
    for i in range(n):
        move[j][i] = move[j - 1][move[j - 1][i]]
a = list(map(int, input().split()))
empty_desks = 0
student_replaced = [0] * n
for i in move[0]:
    student_replaced[i] = 1
for i in range(n):
    if not student_replaced[i]:
        empty_desks += 1
lessons = (max(a) - n) // empty_desks
b = [0] * n
students = set_with_geq(n)
for end_student in a:
    if end_student <= n:
        students.delete(end_student)
seen = [0] * (n + 1)
for start_desk in range(n):
    end_desk = start_desk
    for j in range(30):
        if lessons >> j & 1:
            end_desk = move[j][end_desk]
    if not seen[end_desk]:
        seen[end_desk] = 1
        b[start_desk] = a[end_desk]
    else:
        b[start_desk] = students.geq(a[end_desk])
        students.delete(b[start_desk])
print(' '.join(map(str, b)))