AtCoder Beginner Contest 336 F,G問題メモ

from collections import deque

import numpy as np


def rot(a, k):
    b = a.copy()
    if k == 0:
        b[:-1, :-1] = a[:-1, :-1][::-1, ::-1]
    elif k == 1:
        b[:-1, 1:] = a[:-1, 1:][::-1, ::-1]
    elif k == 2:
        b[1:, :-1] = a[1:, :-1][::-1, ::-1]
    elif k == 3:
        b[1:, 1:] = a[1:, 1:][::-1, ::-1]
    return b


h, w = map(int, input().split())

s = np.arange(1, h * w + 1, dtype=np.int8).reshape((h, w))
g = np.array([list(map(int, input().split())) for _ in range(h)], dtype=np.int8)

if (s == g).all():
    print(0)
    exit()

q = deque([(0, s)])
fwd = {s.data.tobytes(): 0}

while q:
    c, a = q.popleft()
    if c >= 10:
        continue
    for r in range(4):
        b = rot(a, r)
        key = b.data.tobytes()
        if key in fwd:
            continue
        fwd[key] = c + 1
        q.append((c + 1, b))

q = deque([(0, g)])
bwd = {g.data.tobytes(): 0}
INF = 1 << 60
ans = INF

while q:
    c, a = q.popleft()
    key = a.data.tobytes()
    if key in fwd:
        ans = min(ans, c + fwd[key])
        continue
    if c >= 10:
        continue
    for r in range(4):
        b = rot(a, r)
        key = b.data.tobytes()
        if key in bwd:
            continue
        bwd[key] = c + 1
        q.append((c + 1, b))

if ans == INF:
    print(-1)
else:
    print(ans)

細かいことだけど、Numpyのboolean-indexで ATrue, False, True としたら $A$ の $0,2$ 行目が抽出された新たな配列ができるけど、 これを2次元にした場合 ATrue, False, True], [False, True, True は $[A[0,1], A[2,2]]$ という「対角成分の1次元配列」になって（Trueの数が合わないとそもそもエラー）、それぞれがTrueの行・列を抜き出した2次元配列にはなってくれないんだね。

import numpy as np
from numba import njit


class UnionFind:
    def __init__(self, n):
        self.table = [-1] * n

    def root(self, x):
        stack = []
        tbl = self.table
        while tbl[x] >= 0:
            stack.append(x)
            x = tbl[x]
        for y in stack:
            tbl[y] = x
        return x

    def find(self, x, y):
        return self.root(x) == self.root(y)

    def unite(self, x, y):
        r1 = self.root(x)
        r2 = self.root(y)
        if r1 == r2:
            return False
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            self.table[r1] += d2
        else:
            self.table[r1] = r2
            self.table[r2] += d1
        return True

    def get_size(self, x):
        return -self.table[self.root(x)]


def precompute_factorials(n, MOD):
    f = 1
    factorials = [1]
    for m in range(1, n + 1):
        f = f * m % MOD
        factorials.append(f)
    f = pow(f, MOD - 2, MOD)
    finvs = [1] * (n + 1)
    finvs[n] = f
    for m in range(n, 1, -1):
        f = f * m % MOD
        finvs[m - 1] = f
    return factorials, finvs


@njit('(i8[:,:],i8)', cache=True)
def determinant(matrix, MOD):
    def mod_pow(x, a, MOD):
        ret = 1
        cur = x
        while a > 0:
            if a & 1:
                ret = ret * cur % MOD
            cur = cur * cur % MOD
            a >>= 1
        return ret

    n, m = matrix.shape
    assert n == m

    det = 1
    for i in range(n - 1):
        if matrix[i, i] == 0:
            for k in range(i + 1, n):
                if matrix[k, i] != 0:
                    matrix[i], matrix[k] = matrix[k], matrix[i]
                    det *= -1
                    det %= MOD
                    break
            else:
                return 0
        inv = mod_pow(matrix[i][i], MOD - 2, MOD)
        for k in range(i + 1, n):
            if matrix[k, i] != 0:
                r = matrix[k, i] * inv % MOD
                for j in range(i, n):
                    matrix[k, j] -= r * matrix[i, j]
                    matrix[k, j] %= MOD
        det *= matrix[i, i]
        det %= MOD
    return det


def solve(xxx):
    in__deg = [0] * 8
    out_deg = [0] * 8
    matrix = np.zeros((8, 8), np.int64)
    MOD = 998244353

    # コーナー処理
    if sum(xxx[1:]) == 0:  # X0000 のみ
        return 1
    if sum(xxx[:-1]) == 0:  # X1111 のみ
        return 1
    if sum(xxx[1:-1]) == 0:  # X0000 と X1111 が両方あり、他が無い
        return 0

    uft = UnionFind(8)
    use_vertices = set()

    for i in range(16):
        x = xxx[i]
        if x == 0:
            continue
        u = i >> 1
        v = i & 0b111
        out_deg[u] += x
        in__deg[v] += x
        matrix[u, v] -= x
        uft.unite(u, v)
        use_vertices.add(u)
        use_vertices.add(v)

    # 連結性チェック
    leaders = {uft.root(i) for i in use_vertices}
    if len(leaders) != 1:
        return 0

    # 自明な始点と終点があるか調べる
    s = -1
    g = -1
    for i in range(8):
        if in__deg[i] == out_deg[i]:
            continue
        if in__deg[i] + 1 == out_deg[i]:
            if s != -1:
                return 0
            s = i
        elif in__deg[i] == out_deg[i] + 1:
            if g != -1:
                return 0
            g = i
        else:
            return 0

    for i in range(8):
        matrix[i, i] += out_deg[i]

    facts, finvs = precompute_factorials(1_000_000, MOD)

    if s == g == -1:
        # 全ての入次数＝出次数 なので、最初と最後が同じ並びのものを全探索
        # matrix に関しては、どの頂点を始点・終点と仮定しても自己ループ辺が追加されるだけなので変化が無い＝1回だけ計算して値を使い回せる
        is_zero_row = np.all(matrix == 0, axis=1)
        is_zero_col = np.all(matrix == 0, axis=0)
        assert (is_zero_row == is_zero_col).all()
        matrix = matrix[~is_zero_row][:, ~is_zero_col]
        det = int(determinant(matrix, MOD))

        ans = 0
        for g in use_vertices:
            tmp = 1
            for j in use_vertices:
                od = out_deg[j] + 1 if j == g else out_deg[j]
                tmp *= facts[max(1, od - 1)]
                tmp %= MOD
            ans += det * tmp % MOD
        ans %= MOD
    else:
        # 始点と終点は決まっている
        matrix[g, s] -= 1
        matrix[g, g] += 1

        is_zero_row = np.all(matrix == 0, axis=1)
        is_zero_col = np.all(matrix == 0, axis=0)
        assert (is_zero_row == is_zero_col).all()
        matrix = matrix[~is_zero_row][:, ~is_zero_col]
        det = int(determinant(matrix, MOD))
        tmp = 1
        for j in use_vertices:
            od = out_deg[j] + 1 if j == g else out_deg[j]
            tmp *= facts[max(0, od - 1)]
            tmp %= MOD
        ans = det * tmp % MOD

    # 辺は区別しないので並び方の分を減らす
    for x in xxx:
        ans *= finvs[x]
        ans %= MOD

    return ans


xxx = list(map(int, input().split()))
ans = solve(xxx)
print(ans)