AtCoder Beginner Contest 120

全ACできた。

atcoder.jp

初めは01と10の探索から始めたが、発見をした後、2個進める必要があるためforじゃ回しにくいと、 そもそもindexの扱いにして最初と最後のindexの操作を実装した、、が1文字目のindex走査がかなり複雑になり、 簡単なテストケースは通ったものの、一回WA。

300点でこのindex操作量はなかなかないなと思いつつ、テストケースを作ってるうちに気が付く。 これ、どのような01の順序でも、消えない順序がない。 つまり、0と1がどのような順番で存在しても、必ず消えるだけ消えることに気が付く。 ということで0か1のどちらか少ないほうの倍がそのまま答えになる。

def solve(S):
    return min(S.count("1"), S.count("0")) * 2

assert (solve("1") == 0)
assert (solve("0") == 0)
assert (solve("10") == 2)
assert (solve("101") == 2)
assert (solve("100") == 2)
assert (solve("01") == 2)
assert (solve("011") == 2)
assert (solve("0110") == 4)
assert (solve("1100") == 4)
assert (solve("111000") == 6)
assert (solve("101010") == 6)
assert (solve("000000") == 0)
assert (solve("010000") == 2)
assert (solve("011000") == 4)
assert (solve("011010") == 6)
assert (solve("011110") == 4)
assert (solve("000011110000") == 8)
assert (solve("101110") == 4)
assert (solve("100010") == 4)
assert (solve("010") == 2)
assert (solve("0" * 10000) == 0)
assert (solve("1" * 10000) == 0)
assert (solve("1" * 5000 + "0" * 1) == 2)
assert (solve("10" * 5000) == 10000)

if __name__ == "__main__":
    S = input()
    print(solve(S))

直前でWAした実装で苦労したのでテストが多い。。

atcoder.jp

まず、"崩壊"は考えにくいので逆順にして"建設"をしていく。

あるa, bをつなぐ橋を建設したとき、 a, b がすでにつながっているならば不便さは変わらず、 a, bがつながっていないならば不便さが変わる。

a, b がつながっているかどうかの判定はUnionFindの出番。

さらに不便さの計算は、K個、L個、P個の島々のセットがあった時、 「自分の島々の個数」*「自分の島々から到達できない島の個数」 であるため以下3つを足し、最後に2で割ればよい。

K * (N-K)
L * (N-L) 
P * (N-P)

素直に実装すると、N個の全島群から不便さをM回計算することになる。島群のため、橋が建設されるたびに計算量は減っていくが、 減る量は特に最初のころはあまり期待できないなのでまだ改善が必要。

そこで、不便さの変化量を考える。 橋を建設するとき、島aが所属する島々の数と島bが所属する島々の数だけに依存し、それ以外が変わることはない。

このため、島aが所属する島々と、島bが所属する島々の不便さをいったん取り除き、 マージした後にできた新しい島々の不便さを再計算してやることで、再度すべての計算をしなくて済む。

初期状態はすべての橋が崩壊しているとき(=一つも建設されていない時)で、この時、不便さはN * (N - 1) // 2。

import sys
sys.setrecursionlimit(100000)

class UnionFind:

    def __init__(self, n):
        self.parents = [i for i in range(n + 1)]

    def root(self, i):
        if self.parents[i] == i:
            return i
        else:
            self.parents[i] = self.root(self.parents[i])
            return self.parents[i]

    def unite(self, i, j):
        self.parents[self.root(self.parents[i])] = self.root(j)

    def is_unite(self, i, j):
        return self.root(i) == self.root(j)


from collections import Counter


def solve(N, M, ABs):
    uf = UnionFind(N)
    counter = Counter([i + 1 for i in range(N)])
    ans = []
    tans = N * (N - 1)
    for a, b in ABs[::-1]:
        ans.append(str(tans // 2))

        if not uf.is_unite(a, b):
            ar = uf.root(a)
            br = uf.root(b)

            tans -= counter[ar] * (N - counter[ar])
            tans -= counter[br] * (N - counter[br])
            counter[br] += counter[ar]
            tans += counter[br] * (N - counter[br])

            uf.unite(a, b)
    return "\n".join(ans[::-1])


assert (solve(2, 1, [(1, 2)]) == "1")
assert (solve(4, 5, [(1, 2), (3, 4), (1, 3), (2, 3), (1, 4)]) == "0\n0\n4\n5\n6")
assert (solve(6, 5, [(2, 3), (1, 2), (5, 6), (3, 4), (4, 5)]) == "8\n9\n12\n14\n15")

if __name__ == "__main__":
    N, M = tuple(map(int, input().split(" ")))
    ABs = [tuple(map(int, input().split(" "))) for _ in range(M)]
    print(solve(N, M, ABs))

AC取れたのはよいが、今回は時間かかりすぎた感が否めない。

あとpythonの再帰回数にやられて謎のRE状況から10分悩んだ+15分ペナルティは痛い。 UnionFindで1000回以上して再帰してREってのはすぐに思いつかなかったわ。。