Jeopardy of Dropped Balls (Codeforces 1575J)
Tags: Implement, dsu

Mr. Chanek has a new game called Dropping Balls. Initially, Mr. Chanek has a grid $$$a$$$ of size $$$n \times m$$$Each cell $$$(x,y)$$$ contains an integer $$$a_{x,y}$$$ denoting the direction of how the ball will move.  $$$a_{x,y}=1$$$ — the ball will move to the right (the next cell is $$$(x, y + 1)$$$);  $$$a_{x,y}=2$$$ — the ball will move to the bottom (the next cell is $$$(x + 1, y)$$$);  $$$a_{x,y}=3$$$ — the ball will move to the left (the next cell is $$$(x, y - 1)$$$). Every time a ball leaves a cell $$$(x,y)$$$, the integer $$$a_{x,y}$$$ will change to $$$2$$$. Mr. Chanek will drop $$$k$$$ balls sequentially, each starting from the first row, and on the $$$c_1, c_2, \dots, c_k$$$-th ($$$1 \leq c_i \leq m$$$) columns.Determine in which column each ball will end up in (position of the ball after leaving the grid).

Input: The first line contains three integers $$$n$$$, $$$m$$$, and $$$k$$$ ($$$1 \leq n, m \leq 1000$$$, $$$1 \leq k \leq 10^5$$$) — the size of the grid and the number of balls dropped by Mr. Chanek.The $$$i$$$-th of the next $$$n$$$ lines contains $$$m$$$ integers $$$a_{i,1},a_{i,2},\ldots,a_{i,m}$$$ ($$$1 \leq a_{i,j} \leq 3$$$). It will satisfy $$$a_{i, 1} \ne 3$$$ and $$$a_{i, m} \ne 1$$$.The next line contains $$$k$$$ integers $$$c_1, c_2, \ldots, c_k$$$ ($$$1 \leq c_i \leq m$$$) — the balls' column positions dropped by Mr. Chanek sequentially.

Output: Output $$$k$$$ integers — the $$$i$$$-th integer denoting the column where the $$$i$$$-th ball will end.

Note: In the first example, the first ball will drop as follows. Note that the cell $$$(1, 1)$$$ will change direction to the bottom direction.  The second and third balls will drop as follows.   All balls will be dropped from the first row and on the $$$c_1, c_2, \dots, c_k$$$-th columns respectively. A ball will stop dropping once it leaves the grid.

Sample Inputs:
5 5 3
1 2 3 3 3
2 2 2 2 2
2 2 2 2 2
2 2 2 2 2
2 2 2 2 2
1 2 1

1 2 2
1 3
1 2

Sample Outputs:
2 2 1

1 2

Solution (Python):
import os
import sys
from io import BytesIO, IOBase
BUFSIZE = 8192

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = 'x' in file.mode or 'r' not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b'\n') + (not b)
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            (self.buffer.truncate(0), self.buffer.seek(0))

class IOWrapper(IOBase):

    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode('ascii'))
        self.read = lambda : self.buffer.read().decode('ascii')
        self.readline = lambda : self.buffer.readline().decode('ascii')
(sys.stdin, sys.stdout) = (IOWrapper(sys.stdin), IOWrapper(sys.stdout))
input = lambda : sys.stdin.readline().rstrip('\r\n')
from math import ceil

def prod(a, mod=10 ** 9 + 7):
    ans = 1
    for each in a:
        ans = ans * each % mod
    return ans

def gcd(x, y):
    while y:
        (x, y) = (y, x % y)
    return x

def lcm(a, b):
    return a * b // gcd(a, b)

def binary(x, length=16):
    y = bin(x)[2:]
    return y if len(y) >= length else '0' * (length - len(y)) + y

class DisjointSetUnion:

    def __init__(self, n):
        self.parent = [*range(n + 1)]
        self.size = [1] * (n + 1)
        (self.min, self.max) = ([*range(n + 1)], [*range(n + 1)])
        self.count = n

    def get(self, a):
        """"""
        if self.parent[a] == a:
            return a
        x = a
        while a != self.parent[a]:
            a = self.parent[a]
        while x != self.parent[x]:
            (self.parent[x], x) = (a, self.parent[x])
        return a

    def union(self, a, b):
        """"""
        (a, b) = (self.get(a), self.get(b))
        if a != b:
            if self.size[a] > self.size[b]:
                (a, b) = (b, a)
            self.parent[a] = b
            self.size[b] += self.size[a]
            self.min[b] = min(self.min[a], self.min[b])
            self.max[b] = max(self.max[a], self.max[b])
            self.count -= 1

    def count_sets(self):
        """"""
        return self.count
for _ in range(int(input()) if not True else 1):
    (n, m, k) = map(int, input().split())
    dsu = DisjointSetUnion(n * m + m)
    a = []
    for i in range(n):
        a += [list(map(int, input().split()))]
    columns = list(map(int, input().split()))
    ans = []
    for c in columns:
        (x, y) = (0, c - 1)
        while 1:
            val = x * m + y + 1
            val = dsu.max[dsu.get(val)]
            (x, y) = ((val - 1) // m, (val - 1) % m)
            if x >= n:
                break
            dsu.union(val, val + m)
            if a[x][y] == 1:
                y += 1
            elif a[x][y] == 3:
                y -= 1
            else:
                x += 1
        ans += [y + 1]
    print(*ans)

Solution (Python):
import sys
sys.setrecursionlimit(100000)

def _r():
    return sys.stdin.buffer.readline()

def rs():
    return _r().decode('ascii').strip()

def rn():
    return int(_r())

def rnt():
    return map(int, _r().split())

def rnl():
    return list(rnt())

def solve(n, m, k, mat, q):
    (sol, cache) = ([], [0] * m)
    for c in q:
        cc = c - 1
        if cache[cc] == 0:
            (r, c, alldown) = (0, c - 1, True)
            while r < n:
                if mat[r][c] == 2:
                    r += 1
                else:
                    alldown = False
                    (mat[r][c], c) = (2, c + 2 - mat[r][c])
            if alldown:
                cache[cc] = c + 1
            sol.append(c + 1)
        else:
            sol.append(cache[cc])
    return sol
(n, m, k) = rnt()
mat = []
for _ in range(n):
    mat.append(rnl())
q = rnl()
print(' '.join(map(str, solve(n, m, k, mat, q))))

Solution (Python):
from bisect import bisect_right, bisect_left, insort

def solve():
    (n, m, k) = map(int, input().split())
    arr = []
    for _ in range(n):
        a = list(map(int, input().split()))
        arr.append(a)
    cols = [[] for _ in range(m)]
    for i in range(n):
        for j in range(m):
            if arr[i][j] != 2:
                cols[j].append(i)
    queries = list(map(int, input().split()))
    for startPos in queries:
        (x, y) = (0, startPos - 1)
        while x != n:
            if arr[x][y] == 2:
                index = bisect_right(cols[y], x)
                if index == len(cols[y]):
                    break
                else:
                    x = cols[y][index]
            elif arr[x][y] == 1:
                arr[x][y] = 2
                cols[y].remove(x)
                y += 1
            else:
                arr[x][y] = 2
                cols[y].remove(x)
                y -= 1
        print(y + 1, end=' ')
solve()