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

[AtCoder Beginners Contest 333 (ABC333)](https://atcoder.jp/contests/abc333) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。


# A - Three Threes 

## 問題
数字 $N$ ($1\leq N \leq 9$) を $N$ 個つなげて得られる文字列を出力してください。

https://atcoder.jp/contests/abc333/tasks/abc333_a

## 考察
同じ文字列をたくさん連結したいとき、例えば "ABABABAB"という文字列がほしいとき、`S="AB"*4` と書くと変数 $S$ に "ABABABAB" が代入されます。

この書き方をつかうと、数字 $N$ を $N$ 個つなげて得られる文字列を出力できます。

## コード
```Python
N = int(input())

str_N = str(N)  # 数字Nを文字列に変換したもの
ans = str_N * N
print(ans)
```
# B - Pentagon

## 問題
正五角形ABCDEがあります。
入力で与えられる $2$ つの線分の長さは等しいですか？

https://atcoder.jp/contests/abc333/tasks/abc333_b

## 考察
じっくり観察してみると、線分の長さは $2$ 種類しかないことが分かります。
$2$ 種類の長さのうち短い方は、次の $10$ 通りしかありません。
```Python
dist1_set = {"AB", "AE", "BA", "BC", "CB", "CD", "DC", "DE", "EA", "ED"}
```

逆にこのsetの中にない線分はすべて長い方の線分になります。

$2$ つの線分 $S,T$ がどちらも短い線分のときと、どちらも長い線分のときは Yes を出力し、それ以外は No を出力すればいいです。

## コード
```Python
S = input()
T = input()

dist1_set = {"AB", "AE", "BA", "BC", "CB", "CD", "DC", "DE", "EA", "ED"}
S_is_dist1 = S in dist1_set
T_is_dist1 = T in dist1_set
if S_is_dist1 == T_is_dist1:
    print("Yes")
else:
    print("No")
```
# C - Repunit Trio

## 問題
レピュニット数(すべての桁の数が $1$ である整数)を $3$ つ足してできる整数のうち、 $N$ 番目に小さい数はいくつですか？

https://atcoder.jp/contests/abc333/tasks/abc333_c

## 考察
制約を見てみると、 $1\leq N \leq 333$ です。そこそこ小さめ。 小さいものから順に数えて $333$ 番目までの数しか聞かれないなら、テクニカルなことをしなくても愚直な全探索が通ったりしそうです。

次にサンプルを見てみます。
サンプル3を見ると、入力が $333$ のとき、出力は $112222222233$ ( $12$ 桁)です。
つまり、 $13$ 桁以上のレピュニット数をつかわなくていいことが分かります。

「$1$ 桁～ $12$ 桁のレピュニット数の中から重複を許して $3$ つ選んで足した値のうち、小さい方から $N$ 番目の数はいくつですか？」という問題に置き換えられます(高校数学では「重複組み合わせ」とよばれています)。

これはfor文を3重に書けば下のコードのように実装できます。

## コード
```Python
def repunit(length):
    return int('1' * length)


repunit3_lst = set()
for i in range(1, 13):
    for j in range(1, 13):
        for k in range(1, 13):
            value = repunit(i) + repunit(j) + repunit(k)
            repunit3_lst.add(value)

N = int(input())
print(sorted(repunit3_lst)[N - 1])
```

## 別解
<details><summary>itertools.combinations_with_replacementをつかった解法</summary>

`itertools.combinations_with_replacement` をつかうと、重複組み合わせの全パターンを列挙できます。

```Python
from itertools import combinations_with_replacement

repunit3_lst = []
for comb in combinations_with_replacement(range(1, 13), 3):
    val = 0
    for c in comb:
        val += int('1' * c)
    repunit3_lst.append(val)

N = int(input())
print(sorted(repunit3_lst)[N - 1])
```
</details>



# D - Erase Leaves

## 問題
$N$ 頂点の木があります。「木から葉を $1$ つ取り除く」と言う操作を何回でもできます。頂点 $1$ を取り除くまでに必要な最小の操作回数を求めてください。

https://atcoder.jp/contests/abc333/tasks/abc333_d

## 考察
木を下図のように分けます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/3e1bd442-8b88-1ef2-ecb3-b6331498cd2a.png)

頂点数 $c$ の木があったとき、その木の頂点全てを削除するためには $c$ 回操作しないといけません。
上の図であれば、サイズ $2$ の部分木とサイズ $3$ の部分木に対してそれぞれ $2$ 回、 $3$ 回だけ操作することで、この $2$ つの部分木は完全に消えます。
そうすると、頂点 $1$ は葉になり、頂点 $1$ を削除することができます。

部分木のサイズたちのうち、最も大きなサイズを $1$ つだけ残してそれ以外を削除し、最後に頂点 $1$ を削除する。という流れになります。

下のコードでは部分木のサイズをDFSで求めていますが、頂点 $1$ を含む辺だけUnion-Findクラスに与えないようにすることでUnion-Findのsize()メソッドでも求められます(好きなのを選んでいいです)。

## コード
```Python
import sys
import pypyjit

sys.setrecursionlimit(10 ** 7)
pypyjit.set_param('max_unroll_recursion=-1')


def size(v, par):
    result = 1
    for ch in G[v]:
        if ch == par:
            continue
        result += size(ch, v)
    return result


N = int(input())
G = [[] for _ in range(N)]

for _ in range(N - 1):
    u, v = map(lambda x: int(x) - 1, input().split())
    G[u].append(v)
    G[v].append(u)

P = []
for v in G[0]:
    P.append(size(v, 0))
ans = sum(P) - max(P) + 1
print(ans)
```
# E - Takahashi Quest

## 問題
以下の $2$ 種類のクエリが $N$ 個とんできます。
- モンスター $x_i$ を倒せるポーションをゲットできる(放置してもok)
- モンスター $x_i$ とバトルする

モンスターとのバトルにすべて勝つことができますか？
勝てる場合は、冒険の途中で持っているポーションの個数の最大値と、ポーションをゲットするタイミングがいつかを教えてください。

https://atcoder.jp/contests/abc333/tasks/abc333_e

## 考察
クエリが来た順に処理するのもいいですが、競プロあるあるの「クエリを逆順に見てみる」をやってみます。

クエリを逆順に見て、まずモンスター $x$ とのバトルが発生したとします。
ポーション $x$ が $1$ つ必要なので、(クエリを逆順に見て)このモンスターとのバトルの後に発生する「ポーション $x$ をゲットできる」のイベントでポーション $x$ をゲットすればいいです。

- $portion =持っているポーションの個数$
- $dic[x]=倒せるのが確定していないモンスターxの数$ 

の2つを管理しながら進めます。

まず処理が簡単なのは、モンスターとバトルするクエリのときです。
ここでポーションを $1$ つ使うはずなので、持っていたはずのポーションの数は $1$ つ増えます。

```Python
if t == 2:
    dic[x] += 1
    portion += 1
```

ポーション $x$ をゲットできるクエリのとき、もしここまでクエリを逆順に見てきてモンスター $x$ とのバトルが発生していなければ、そのポーションを拾う必要はありません。

逆に、モンスター $x$ とのバトルが発生していたとき、つまり `dic[x] > 0` のとき、このクエリでポーション $x$ を拾います。
```Python
if t == 1:
    if dic[x] > 0:
        dic[x] -= 1
        portion -= 1
```

これをまとめると、下のコードになります。
基本的には上に書いたコードをそのまま使いながら、「ポーションをつかったかどうか」も ans_que に記録しています。

## コード
```Python
from collections import defaultdict, deque

N = int(input())
query = [tuple(map(int, input().split())) for _ in range(N)]

dic = defaultdict(int) 
portion = 0  # 持っているポーションの個数
ans = 0  # 持っているポーションの個数の最大値
ans_que = deque()
while query:
    t, x = query.pop()
    if t == 1:
        if dic[x] > 0:
            dic[x] -= 1
            portion -= 1
            ans_que.appendleft(1)
        else:
            ans_que.appendleft(0)
    elif t == 2:
        dic[x] += 1
        portion += 1
    ans = max(ans, portion)

if max(dic.values()) > 0:
    print(-1)
else:
    print(ans)
    print(*ans_que)
```
# F - Bomb Game 2

## 問題
$N$ 人が一列に並んでいます。
$1/2$ ずつの確率で、以下のイベントが発生します。
- 列の先頭の人がいなくなる
- 列の先頭の人が、列の最後尾に移動する

人 $1,2,\cdots ,N$ それぞれについて、人 $i$ が最後まで列に並んでいる $1$ 人になる確率を $\text {mod} 998244353$ で求めてください。

https://atcoder.jp/contests/abc333/tasks/abc333_f

## 考察
いろいろな解き方がありますが、ここでは確率漸化式っぽく見て解いていきます。

$N=3$ のときを考えてみます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/66165620-8ca0-242f-ff0f-bce69307e9ac.png)


最後の $1$ 人になる確率を、列の先頭から順に $a_0, a_1, a_2$ とします。

まずは人 $0$ の生き残る確率 $a_0$ について考えます。
 $1/2$ の確率で列の最後尾に行きます。列の最後尾の人が生き残る確率は $a_2$ です。
 残り $1/2$ の確率で列からいなくなります。生き残る確率をあえて書くならば $0$ です。
 まとめると、こうなります。
 $$a_0 = \frac{1}{2}a_2 + \frac{1}{2}\times 0 \tag{1}$$

 次に人 $1$ の生き残る確率 $a_1$ について考えます。
 $1/2$ の確率で列の先頭に行きます。列の先頭の人が生き残る確率は $a_0$ です。
 $1/2$ の確率で、残り $2$ 人の列の先頭に行きます。
 $N=2$ のときの問題を解かないといけなくなりました。ここではいったんそのときの問題の答えは出ているとして、 $N=2$ のときに最後の $1$ 人になっている確率を列の先頭から $b_0, b_1$ とします。
つまり、$1/2$ の確率で、残り $2$ 人の列の先頭に行くことになり、残り $2$ 人の列の先頭が生き残る確率は $b_0$ です。
まとめると、こうなります。
$$a_1 = \frac{1}{2}a_0 + \frac{1}{2}b_0 \tag{2}$$

同様にして、人 $2$ の生き残る確率 $a_2$ について考えると、下の式が出てきます。
$$a_2 = \frac{1}{2}a_1 + \frac{1}{2}b_1 \tag{3}$$

式 $1$ に式 $3$ を代入して、さらにそこに式 $2$ を代入すると、次の式になります。
$$a_0 = \frac{1}{8}a_0+\frac{1}{8}b_0+\frac{1}{4}b_1$$

すでに $N=2$ の問題が解けていて、 $b_0, b_1$ の値が分かっていれば、 $a_0$ の値が分かります。
また、式 $2,3$ をつかえば、 $a_1, a_2$ の値も分かります。

人数が $N$ 人のときの問題は、人数が $N-1$ 人のときの問題の答えが分かっていれば求められることが分かりました。
人数が $1$ 人のときの問題は、そもそも最初から $1$ 人しかいないので、その人が最後の $1$ 人になる確率は $1$ です。
これを再帰関数で実装すれACできます。

## コード
```Python
import sys
import pypyjit

sys.setrecursionlimit(10 ** 7)
pypyjit.set_param('max_unroll_recursion=-1')

MOD = 998244353
INV = (MOD + 1) // 2  # 1/2 mod 998244353


def solve(x):
    if x == 1:
        return [1]

    prev = solve(x - 1)
    a0 = 0
    for i in range(x - 1):
        a0 += prev[x - i - 2] * inv[i + 2] % MOD
        a0 %= MOD
    a0 = a0 * rui[x] % MOD * pow(rui[x] - 1, MOD - 2, MOD) % MOD
    result = [a0]
    for i in range(1, x):
        result.append((prev[i - 1] + result[-1]) * INV % MOD)
    return result


N = int(input())

rui = [1]  # rui[x]: 2^x mod 998244353
inv = [1]  # inv[x]: (1/2)^x mod 998244353
for i in range(N + 2):
    rui.append(rui[-1] * 2 % MOD)
    inv.append(inv[-1] * INV % MOD)

print(*solve(N))
```