ABC329をPythonで解いてみたよ。(A～F問題)

AtCoder Beginners Contest 329 (ABC329) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。

A - Spread

問題

1文字ずつ半角スペースを入れてね。

考察

Pythonでは、リストの先頭に * をつけて出力すると、要素と要素の間に1つ半角スペースが入った状態で出力されます。(他の問題でもよくつかう出力方法です！)
入力で受け取った文字列をリストにして、これを利用します。

コード

S = list(input())
print(*S)

B - Next

問題

最大の整数の次に大きな整数はいくつですか？

考察

この問題では、同じ数が何個あっても、数がどの位置にあっても、答えは変わりません。
たとえば、 $[3,1,4,1,5]$ と $[3,1,4,1,5,5]$ の答えは同じですし、 $[3,1,4,1,5]$ と $[1,1,3,4,5]$ の答えも同じです。
こういうときは、後で簡単になるように、$A$ はリストではなくセットで受け取ります。

セットから特定の要素 $el$ を取り出すときは、 .discard(el) をつかいます。ちなみに、この操作は $O(1)$ でとても軽いです。

コード

N = int(input())
A = set(map(int, input().split()))
A.discard(max(A))
print(max(A))

C - Count xxx

問題

同じ文字だけをつかった部分文字列は、いくつありますか？

考察

それぞれの文字について、最長の部分文字列の長さを求めます。
たとえば "aaabaabb" は、"a"については最長で $3$ 、"b"については最長で $2$ です。
"a"については最長で $3$ というのは、言い換えると "a", "aa", "aaa" の $3$ 種類の部分文字列をつくることができるということです。
なので答えは、各文字についての最長の部分文字列の長さの総和になります(上の例なら $3+2=5$ が答えです)。

"aaabaabb" を $(a \times 3) + (b \times 1) + (a \times 2)+(b \times 2)$ と表現するのは「ランレングス圧縮」とよばれていて、Pythonでは itertools.groupby をつかうと便利です。

コード

from itertools import groupby
from collections import defaultdict

N = int(input())
S = input()

dic = defaultdict(int)
for el, group in groupby(S):
    dic[el] = max(dic[el], len(list(group)))
print(sum(dic.values()))

D - Election Quick Report

問題

考察

まずは、愚直に解けないかどうかを考えます。

たとえば、こんな解法が思いつきます。

1. (得票数, 人番号) のタプルを、リストか何かに $N$ 人分つっこみます。
2. 特定の人番号のタプルを取り出して、その得票数を $+1$ して元に戻します。
3. そのリストをソートして、最も得票数の多い人の番号を求めます。(得票数が同じときは、人番号の小さい方が選ばれるようにソートします)。

リストを毎回ソートするわけにもいかないですし、特定の人番号がどの位置にあるのかもよく分かりません。
ですが、このリストをSortedSetにすると、その 2つの問題はすべて解決します。

操作の回数は $N$ 回、操作1回あたり計算量は $O(\sqrt N)$ なので、全体で $O(N \sqrt N)$ で解けます。

コード

'''
tatyamさん作の、SortedSetです。
使わせていただき、ありがとうございます！
https://github.com/tatyam-prime/SortedSet/blob/main/SortedSet.py

・使い方(個人的まとめ)
s=SortedSet()
s.a: SortedSetの中身を返す。
len(s), x in s, x not in s: リストと同じ要領で使える。
s.add(x): xを追加してTrueを返す。ただしxがすでにs内にある場合、xは追加せずにFalseを返す。
s.discard(x): xを削除してTrueを返す。ただしxがs内にない場合、何もせずにFalseを返す。
s.lt(x): xより小さい最大の要素を返す。もし存在しないなら、Noneを返す。
s.le(x): x　以下の　最大の要素を返す。もし存在しないなら、Noneを返す。
s.gt(x): xより大きい最小の要素を返す。もし存在しないなら、Noneを返す。
s.ge(x): x　以上の　最小の要素を返す。もし存在しないなら、Noneを返す。
s.index(x): xより小さい要素の数を返す。
s.index_right(x): x以下の要素の数を返す。

・使い方URL
https://github.com/tatyam-prime/SortedSet
'''
# (SortedSetのコードは省略。上のURLからコピペしてね。)


N, M = map(int, input().split())
A = list(map(lambda x: int(x) - 1, input().split()))

cnt = [0] * N
ss = SortedSet()
for i in range(N):
    ss.add((0, -i))
inf = 1 << 60
for a in A:
    ss.discard((cnt[a], -a))
    cnt[a] += 1
    ss.add((cnt[a], -a))
    _, mi = ss.le((inf, 0))
    print(-mi + 1)

別解

$O(N)$ 解法

投票されたとき、当選しているのは、「今回投票された人」もしくは「前回の時点で1位だった人」 のどちらかです。
この事実を利用すれば、投票されるごとに $O(1)$ で1位の人を答えられて、全体で $O(N)$ で問題を解けます。

N, M = map(int, input().split())
A = list(map(lambda x: int(x) - 1, input().split()))

cnt = [0] * N
prev_winner = 0

for a in A:
    cnt[a] += 1
    if (cnt[prev_winner], -prev_winner) < (cnt[a], -a):
        winner = a
    else:
        winner = prev_winner
    print(winner + 1)
    prev_winner = winner

E - Stamp

問題

文字列 $T$ のスタンプをうまく押すことで、文字列 $S$ はできますか？

考察

競プロあるあるの1つ「逆から考える」をつかいます。
つまり、"##########..." の文字列から文字列 $S$ にするのではなく、文字列 $S$ から"##########..." の文字列にできるかどうかを考えます。

文字列 $S$ の中に 文字列 $T$ と一致している部分があれば、その部分をすべて "#" にすることにします。
また、文字列 $S$ の中に、文字列 $T$ の一部分が "#" になっているような文字列と一致している部分があっても、その部分をすべて "#" にすることができます。(実際の操作の逆の手順をしていることをイメージすると、この操作がokなのが分かりやすいと思います)
これを今回は「操作」とよぶことにします。

これを貪欲に操作していきたいです。
「操作を行える場所はどこかなー？」と毎回左端から右端まで見ていってると時間がかかるので、少し工夫します。

どこか文字列 $S$ の中で操作をしたとします。 $[l,l+M)$ の範囲の文字をすべて "#" にしたとします。
すると、新たに文字列 $S$ の中で操作できるようになる場所は、$[l-M+1, l+1)$ から $[l+M-1,l+2M-1)$ の計 $2M-2$ 箇所だけです。
今回は $M \leq 5$ なので、この特性を利用すれば、全体で $O(NM)$ で問題を解くことができます。
いくつか手法はありますが、下のコードでは、先に操作をできる箇所を $O(N)$ で探索してスタックに入れて、スタックから取り出したインデックスに対して操作を行い、操作するごとに先ほどの $2M-2$ 箇所が操作できないかどうか確認し、操作できるならスタックにそのインデックスを入れています。(下のコードも参考にしてください)

コード

def can_stamp(left_idx: int) -> bool:
    for el1, el2 in zip(S[left_idx:left_idx + M], T):
        if el1 not in ['#', el2]:
            return False
    return True


def stamp(left_idx: int) -> None:
    for dist in range(M):
        i = left_idx + dist
        if i >= N:
            break
        S[i] = '#'


N, M = map(int, input().split())
S = list(input())
T = list(input())

seen = set()
stamp_stack = []
for i in range(N - M + 1):
    if can_stamp(i):
        stamp_stack.append(i)
        seen.add(i)

while stamp_stack:
    idx = stamp_stack.pop()
    stamp(idx)
    for i in range(idx - M + 1, idx + M):
        if not 0 <= i < N - M + 1: continue
        if can_stamp(i) and i not in seen:
            stamp_stack.append(i)
            seen.add(i)
for s in S:
    if s != '#':
        print("No")
        exit()
print("Yes")

F - Colored Ball

問題

「箱 $A$ の中身を、箱 $B$ の中に全部入れる」の操作をたくさんして、その操作ごとに箱 $B$ の中にある玉の色が何種類あるか答えてね。

考察

マージテクをつかいます。
この問題では、箱 $A$ の中身を箱 $B$ にすべて移す操作をたくさんします。
箱 $B$ の中身よりも箱 $A$ の中身の方が多いとき、 箱 $B$ の中身を箱 $A$ にすべて移して、そのあと箱 $A$ と箱 $B$ をひっくり返すことにすると、移す玉の個数が減ることで計算量が落ちて操作1回あたり $O(\log N)$ になります。

玉の個数ではなく種類数のみ聞かれているので、セットの結合になります。
セットの結合は、bit演算と同じ感じでor演算がつかえます(下のコードも参考にしてください)。

コード

N, Q = map(int, input().split())
C = list(map(int, input().split()))

se = [{c} for i, c in enumerate(C)]

for _ in range(Q):
    a, b = map(lambda x: int(x) - 1, input().split())
    if len(se[a]) < len(se[b]):
        se[b] |= se[a]
        se[a] = set()
    else:
        se[a] |= se[b]
        se[a], se[b] = se[b], set()
    print(len(se[b]))