Zenn
😽

[python] ARC107 個人的メモ [Atcoder]

2020/11/01に公開
Python
競技プログラミング
AtCoder
tech

ARC107

所感

BA2完 971パフォ
入緑した
やったぜ

A - Simple Math

最初に見たとき，さっぱりわからなくて焦った
とは言え，B問題の後に見てサンプル1を書いて考えたら割とあっさりわかってよかった
まず総和の公式と言うのがある

\sum_{x=1}^N x = \frac{N(N+1)}{2}

これを使って問題の式を解いていくと以下の様に解ける

\begin{aligned} \sum_{a=1}^A \sum_{b=1}^B \sum_{c=1}^C abc&=\sum_{a=1}^A \sum_{b=1}^B ab\frac{C(C+1)}{2}\\ &=\frac{C(C+1)}{2}\sum_{a=1}^Aa\frac{B(B+1)}{2}\\ &=\frac{C(C+1)}{2}\frac{B(B+1)}{2}\frac{A(A+1)}{2} \end{aligned} 
#!/usr/bin/env python3
A, B, C = map(int, input().split())
MOD = 998244353

c = C * (C + 1) // 2
b = B * (B + 1) // 2
a = A * (A + 1) // 2

print(a * b * c % MOD)

B - Quadruple

abcdをそれぞれで全探索するとO(N^4)でTLE
ということでa-c=ac,d-b=dbとしてac-db=Kとして考えた
これならacを全探索すれば，dbはそれぞれの場合でO(1)で求まるから，全体の計算量はO(N)で済む

1つ目の制約より，dbのとれる範囲には限りがあるので，範囲外になる場合は除く
次にa-c=ac,\ d-b=dbとなるabcdを求める必要がある
ac,\ dbが共に正ならそれぞれN-ac,\ N-dbで求まる
これは，区間1からNに長さac,\ dbの棒が何本入るかを考えれば分かる

acが負の場合はaとcを入れ替えれば正になって同じようにN-acで求まる
dbも同様に考えれば良いから，a-c=ac,\ d-b=dbとなるabcdはそれぞれN-|ac|,\ N-|db|で求まると分かる
で，各acでaとcの取り得る場合とdとbの取り得る場合を掛ければおｋ

公式解説だと，abとcdでセットにしてた

#!/usr/bin/env python3
N, K = map(int, input().split())

ans = 0
for ac in range(1 - N, N):
    db = ac - K
    if db < 1 - N or N - 1 < db:
        continue
    ans += (N - abs(ac)) * (N - abs(db))
print(ans)

C - Shuffle Permutation

wa
2つの操作が独立してそうというのは分かったけど，そこから考察がうまいこと進まなかった
unionfindも全然思いつかなかった

列のswapと行のswapの操作は互いに独立している操作だから，両操作の場合の数の積が出力となればおｋ

制約の3つ目より全ての行列成分が異なると分かる
なので，重複とかは考えなくて良い

まず行のswapについて考える
例えば互いに操作可能な4つの行abcdがあるとする
操作回数に制限はないから，これらの行は好きな順に並べ替えられる
よって，これらの行の組み合わせは4!通りと分かる
で，他の互いに操作可能な行のグループでも同じように組み合わせを求める
それらの積が最終的に得られる全ての行の組み合わせとなる

同様に列でも全ての組み合わせを求め，行の組み合わせとの積を求めれば解が求まる

互いに操作可能なグループを求めるにはunionfindを用いた

#!/usr/bin/env python3
class UnionFind:
    def __init__(self, n: int):
        self.root = [-1] * n

    def find(self, x: int):
        if self.root[x] < 0:
            return x
        self.root[x] = self.find(self.root[x])
        return self.root[x]

    def unite(self, x: int, y: int):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return False
        if self.root[x] > self.root[y]:
            x, y = y, x
        self.root[x] += self.root[y]
        self.root[y] = x
        return True

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

    def size(self, x: int) -> bool:
        return self.root[self.find(x)] * -1


def main():
    from math import factorial

    N, K = map(int, input().split())
    matrix = [list(map(int, input().split())) for _ in range(N)]
    mod = 998244353

    row_group = UnionFind(N)
    line_group = UnionFind(N)

    for i in range(N):
        for j in range(i + 1, N):
            row = True
            line = True
            for k in range(N):
                if matrix[i][k] + matrix[j][k] > K:
                    row = False
                if matrix[k][i] + matrix[k][j] > K:
                    line = False
            if row:
                row_group.unite(i, j)
            if line:
                line_group.unite(i, j)

    ans = 1
    for i in range(N):
        # root[i]<0というのは，root[i]がグラフの木の根になっているということ
        # 要するに同じグループを重複してカウントしないようにしている
        if row_group.root[i] < 0:
            ans *= factorial(row_group.size(i))
        if line_group.root[i] < 0:
            ans *= factorial(line_group.size(i))

    print(ans % mod)


main()

参考資料

Discussion