Submission #17757495

---

Source Code Expand

import sys
import numpy as np
import numba
from numba import njit, b1, i4, i8, f8

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

def from_read(dtype=np.int64):
    return np.fromstring(read().decode(), dtype=dtype, sep=' ')


def from_readline(dtype=np.int64):
    return np.fromstring(readline().decode(), dtype=dtype, sep=' ')

"""def naive(col, n):
    def mex(a,b):
        for i in range(3):
            if a == i or b == i:
                continue
            return i
    dp = np.zeros_like(col)
    x = n
    for i in range(len(col)):
        dp[i] = mex(col[i], x)
        x = dp[i]
    return dp
col = np.random.randint(0, 3, 30)
for _ in range(10):
    print(col)
    x = np.random.randint(0, 3)
    col = naive(col, x)"""

@njit((i8, i8), cache=True)
def mex(a, b):
    for i in range(3):
        if a == i or b == i:
            continue
        return i

@njit
def update_row(row, x):

    dp = np.zeros_like(row)
    for i in range(len(row)):
        dp[i] = mex(row[i], x)
        x = dp[i]
    return dp

@njit
def solve_small(N, A):
    ans = np.zeros(3, np.int64)
    ans[A[0]] += 1
    row = A[1:N]
    col = A[N:]
    for x in row:
        ans[x] += 1
    for x in col:
        ans[x] += 1
    for i in range(N - 1):
        row = update_row(row, col[i])
        for x in row:
            ans[x] += 1
    return ans

@njit((i8, i8[:]), cache=True)
def main(N, A):
    if N <= 3:
        return solve_small(N, A)
    # (1,1) セルは茶番
    ans = np.zeros(3, np.int64)
    ans[A[0]] += 1
    row = A[1:N]
    col = A[N:]
    for x in row:
        ans[x] += 1
    for x in col:
        ans[x] += 1
    # とりあえず、10 回くらい進める
    for i in range(3):
        row = update_row(row, col[i])
        for x in row:
            ans[x] += 1
    col = col[3:]
    K = len(col)
    # のこり K 行の計算がある。
    for i in range(2, N - 1):
        # いま入っている数はそのまま右下に移動する
        # 集計回数は？
        k = (N - 2) - i
        k = min(k, K)
        ans[row[i]] += k
    # 左 3 つの数値のみ計算していく
    row = row[:3]
    for i in range(K):
        row = update_row(row, col[i])
        a, b, c = row
        ans[a] += 1
        ans[b] += 1
        ans[c] += K - i
    return ans

N = int(readline())
A = from_read()

print(*main(N, A))

Submission Info

Judge Result