PythonでABC335を解いてみたよ。(A～F問題) - Qiita (2024)

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

問題

文字列 $S$ のは必ず $\text{2023}$ で終わります。
この $S$ の末尾1文字を $\text{4}$ に変えてください。

考察

ans = "" として、空文字列の変数 $ans$ を用意します。
for文で文字列 $S$ の文字を1つずつ $ans$ にくっつけていきます。ただし、最後の文字だけは $ans$ にくっつけないで、代わりに $\text{4}$ をくっつけます。
この変数 $ans$ が答えになっているので、そのまま出力すると正解です。

コード

S = input()n = len(S)ans = ""for i in range(n - 1): ans += S[i]ans += '4'print(ans)

別解

S = input()ans = S[:-1] + '4'print(ans)

S = list(input())S[-1] = '4'print(*S, sep="")

問題

$x+y+z \leq N$ となるような非負整数 $x,y,z$ の組をすべて出力してね(昇順にならべておいてね)。

考察

$N \leq 21$ です。 $x,y,z$ は非負整数なので、 $0 \leq x,y,z \leq 21$ です。
$x,y,z$ すべてそれぞれ $22$ 通りしかなくて、全パターン試しても $22^3=10648$ 通りだけです。
これなら $x,y,z$ の組を全パターン試して、その総和が $N$ 以下のものだけを選ぶ作戦でもよさそうです。

昇順にならべて出力しないといけないのについては、下のコードのように $x,y,z$ の順番でfor文を回せば昇順に出力されます。

コード

N = int(input())for x in range(N + 1): for y in range(N + 1): for z in range(N + 1): if x + y + z <= N: print(x, y, z)

おまけ

N = int(input())lst = []for x in range(N + 1): for y in range(N + 1): for z in range(N + 1): if x + y + z <= N: lst.append((x, y, z))lst.sort() # 本当は必要ないけど、変に考察の時間をのばすくらいならこの方がはやいfor tpl in lst: print(*tpl)

for x in range(N + 1): for y in range(N + 1): for z in range(N + 1): pass # 処理文

from itertools import productfor x, y, z in product(range(N + 1), repeat=3): pass # 処理文

from itertools import productN = int(input())for pro in product(range(N + 1), repeat=3): if sum(pro) <= N: print(*pro)

問題

龍は上下左右 $1$ マスだけ頭(龍の先頭、パーツ $1$ )を動かすことができ、動かすと残りのパーツ $i$ は動かす前のパーツ $i-1$ の座標に移動します。
龍の動くクエリ・龍のパーツ $i$ の座標を出力するクエリの $2$ 種類が合計 $Q$ 回とんでくるので、がんばって処理してください。

考察1 RLUDの辞書

龍の動くクエリで、動く方向が"RLUD"のいずれか1文字で入力されます。
個の問題だけに限らず、このタイプの入力は、先に辞書をつくっておくとラクなことが多いです。

"""RLUDに対応するベクトルの辞書をつくっておく"""V = {'R': (1, 0), 'L': (-1, 0), 'U': (0, 1), 'D': (0, -1) }

考察2 dequeをつかう

最初、龍のパーツの座標はパーツ $1$ から順番に、 $(1, 0), (2, 0), \cdots , (N, 0)$ です。
龍が動くクエリが来たとき、これは $(新しいパーツ1の座標), (1, 0), (2, 0), \cdots ,(N-1,0)$ になります。左に移動後のパーツ $1$ の座標を入れて、右端にある座標を取り除けばokです。

dequeをつかうと、右端の要素を取り出すのも、左端に要素を加えるのも $O(1)$ でできます。。が、龍のパーツの座標を答えるクエリが $O(N)$ になります。
(dequeの添え字アクセスは $O(1)$ じゃなくて $O(N)$ なので、添え字アクセスでdequeの中央付近を見に行くのはやめたほうがいいです。)
全体で $O(NQ)$ になってしまいTLEしてしまうので、この解法はナシです......

考察3 リストをつかう

龍が動くクエリで右端にある座標を取り除いていましたが、右端は取り除かずにそのまま残して、パーツ $i$ は左から $i$ 番目の座標だとしても大丈夫です。
ということで、龍が動くクエリでは、左端に移動後のパーツ $1$ の座標を入れる操作だけをすることにします。

Pythonのリストは要素を追加するメソッド .append() があります。これはリストの右端に要素をつけ加えるメソッドです。
なので、龍のパーツの座標をパーツ $1$ から順番に、 $(1, 0), (2, 0), \cdots , (N, 0)$ としていましたが、これを逆順に並べて、 $(N, 0), (N-1,0), \cdots ,(1,0)$ とします。こうすると、新しい座標を右端に .append() をつかって入れられます。
また、龍のパーツの座標を答えるクエリでは、パーツ $i$ の座標はリストの末尾から $i$ 番目のものを指しているので、$O(1)$ で答えられます。

コード

V = {'R': (1, 0), 'L': (-1, 0), 'U': (0, 1), 'D': (0, -1) }N, Q = map(int, input().split())lst = [(i, 0) for i in range(1, N + 1)]lst = lst[::-1] # .append()がつかえるように、逆順に並べるfor _ in range(Q): t, q = input().split() if t == '1': vx, vy = V[q] prev_x, prev_y = lst[-1] new_x, new_y = prev_x + vx, prev_y + vy lst.append((new_x, new_y)) else: print(*lst[-int(q)])

問題

考察

次の図のように、左上をスタート地点にして反時計回りに渦巻きのように数字を埋めていく作戦でいきます。

以下、ここでは $N=7$ のときだけ考えて解いていきます。
まず、 $1$ ～ $24$ を埋めることを考えます。
$1$ ～ $24$ をこんな感じで $4$ つに分割して考えます。

$1$ ～ $6$ はこのようなコードで埋められます。

N = int(input())Grid = [[0] * N for _ in range(N)]x = 1ni, nj = 0, 0for _ in range(6): Grid[ni][nj] = x x += 1 ni += 1

次に $7$ ～ $12$ はこのようなコードで埋められます。

for _ in range(6): Grid[ni][nj] = x x += 1 nj += 1

$13$ ～ $18$ と $19$ ～ $24$ についても同様です。

$25$ ～ $40$ も次の図のように $4$ 分割して同じ流れで埋めていきます。

この方針ですべての数字を条件を満たしながら埋めることができます。
詳しい実装は下のコードを参考にしてください。

コード

N = int(input())Grid = [[0] * N for _ in range(N)]Grid[N // 2][N // 2] = 'T'x = 1for i in range(N // 2): ni, nj = i, i length = N - 1 - i * 2 for _ in range(length): Grid[ni][nj] = x x += 1 ni += 1 for _ in range(length): Grid[ni][nj] = x x += 1 nj += 1 for _ in range(length): Grid[ni][nj] = x x += 1 ni -= 1 for _ in range(length): Grid[ni][nj] = x x += 1 nj -= 1for row in Grid: print(*row)

別解

V = [(0, -1), (-1, 0), (1, 0), (0, 1)]N = int(input())Grid = [[None] * N for _ in range(N)]Grid[N // 2][N // 2] = 'T'x = 1i, j = 0, 0while True: Grid[i][j] = x for vi, vj in V: ni, nj = i + vi, j + vj if not (0 <= ni < N and 0 <= nj < N): continue if Grid[ni][nj] is None: i, j = ni, nj x += 1 break # 上下左右どの方向にも行けなかったら数字はすべて埋まっている → このwhile文を終了する。 else: breakfor row in Grid: print(*row)

問題

$N$ 頂点の連結無向グラフがあります。各頂点には正整数が書かれています。
頂点 $1$ から頂点 $N$ までの単純パスで、そのパスの順に頂点に書かれた整数を並べた数列を $S$ とします。
広義単調増加な $S$ のうち、$S$ の要素の種類数の最大値はいくつですか？

考察

数列 $S$ は広義単調増加になっていないといけない(そうでないと得点が強制的に $0$ になる)ので、頂点に書かれた整数が下がるような移動はできないことにしましょう。つまり連結な無向グラフではなく、非連結の可能性もある有向グラフとして見ていきます。

気持ち的には、書かれた整数が小さいものから順番に見ていき、辺 $u \rightarrow v$ があって $A_u < A_v$ のとき、種類数について kind[v] = max(kind[v], kind[u] + 1) としていきたいです。頂点 $v$ について見ているときに、辺 $u \rightarrow v$ があって $A_u < A_v$ となる頂点 $u$ がほしくなります。いったんこれをコードにします。

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

書かれた整数が下がるような移動ができないように有向グラフをつくるので、このグラフはトポロジカルソートできます。
トポロジカルソートしておくと、ソートされた順に頂点を見ていけばよくなります。

また、 $A_u = A_v$ で $u \rightarrow v$ の辺があるとき、$kind[u] = kind[v]$ になります。この有向グラフでは $A_u = A_v$ のときは逆の辺 $v \rightarrow u$ もあるので、頂点 $u, v$ は強連結です。
強連結な頂点たちそれぞれについて kind[v] = max(kind[v], kind[u] + 1) (← この $u, v$ は辺 $u \rightarrow v$ があって $A_u < A_v$ の2頂点)を実行して種類数を求めます。その種類数の最大値を $ma$ としたとき、強連結な頂点たちすべてについて、種類数を $ma$ としてokです。

これらはac-library-pythonにあるsccをつかうと実装できます。

これをまとめると、下のコードになります。

コード

from atcoder.scc import SCCGraphN, M = map(int, input().split())A = list(map(int, input().split()))scc = SCCGraph(N)G = [[] for _ in range(N)]for _ in range(M): u, v = map(lambda x: int(x) - 1, input().split()) if A[u] < A[v]: G[v].append(u) scc.add_edge(u, v) elif A[v] < A[u]: G[u].append(v) scc.add_edge(v, u) elif A[u] == A[v]: scc.add_edge(u, v) scc.add_edge(v, u)kind = [0] * Nkind[0] = 1for group in scc.scc(): for v in group: for nv in G[v]: if kind[nv] > 0: kind[v] = max(kind[v], kind[nv] + 1) ma = max(kind[v] for v in group) for v in group: kind[v] = maprint(kind[N - 1])

問題

$N$ 個のマスと数列 $A=(A_1, A_2, \cdots ,A_N)$ があります。
マス $1$ は黒く塗られていて、それ以外のマスは白色です。駒がマス $1$ に置かれています。
以下の操作を $0$ 回以上行います。

• 駒のいるマスを $i$ として、適当な正整数 $x$ を決めて、駒をマス $i + A_i \times x$ に移動させる。駒の移動先のマスを黒く塗る。

操作を終えた時点での黒く塗られたマスの集合の数を $\text{mod} 998244353$ で求めてください。

考察1 (解いてみる)

左端のマスの番号は本来は $1$ ですが、この考察では0-indexedでマス $0$ からスタートするとします。
左のマス $0$ から順に見ていきます。
black[i] : マス $i$ が黒色に塗られているときのパターン数
white[i] : マス $i$ が白色のときのパターン数
とします。左から順に見ていくので、添え字が $i$ のときはマス $i+1$ 以降はすべて白色だとしていることに注意してください。
言い換えると、マス $0$ からマス $i$ までの $i+1$ 個のマスだけを考えて、マス $i$ が黒色に塗られているときのパターンの数が black[i] です(white[i]も同様です)。
このルールでサンプル1を解いてみます。

マス $0$ は必ず黒色に塗られていると問題に書かれているので、 black[0], white[0] = 1, 0 です。
マス $0$ に書かれている整数は $1$ なので、マス $0$ からマス $1, 2, 3, 4$ を黒色に塗れます。この黒色に塗れるマスたちにblack[0] の値をプラスします。

これで black[1] の値は確定しました(マス $1$ より手前のマスのblack についてはすべて処理済なのが理由です)。
次に white[1] の値を考えます。これは、直前のマスが黒色のパターン数と直前のマスが白色のパターン数の和になります。
ということで、このようになります。

次に マス $1$ に書かれた整数が $2$ なので、マス $3$ を黒色に塗れます。
また、white[2] の値も決められます。
これを処理するとこうなります。

この方針で最後まで進めると、表はこのようになります。

これを実装すると、次のようなコードになります。

MOD = 998244353N = int(input())A = list(map(int, input().split()))black, white = [0] * N, [0] * Nblack[0] = 1for i in range(N - 1): for j in range(i + A[i], N, A[i]): black[j] += black[i] black[j] %= MOD white[i + 1] = (black[i] + white[i]) % MODans = (black[-1] + white[-1]) % MODprint(ans)

しかし、これではTLEしてしまいます。。。

考察2 TLEの原因・その対策

競プロあるあるで、「最悪のケースを考えてみる」というのがあります。
意地悪ですが、例えば $A=[1, 1, 1, 1, \cdots ,1]$ のとき、上のコードを実行するととんでもなく時間がかかります。
悪さをしているのは、上のコードのこの部分です。

 for j in range(i + A[i], N, A[i]): black[j] += black[i] black[j] %= MOD

これだと $O(N^2)$ かかってしまいます。
これの対策としては、black[i+A[i]] += black[i] だけを実行しておき、$i+A[i]$ のインデックスに「この先も $A[i]$ 刻みで $\text{black}[i]$ だけ足し算する」というのを記録させておく方法があります。
つまり、次のようなコードに改変です。

# dp[i][v]: iからv刻みでdp[i][v]だけ足し算するというメモdp = [[0] * MOD for _ in range(N)]for i in range(N - 1): dp[i][A[i]] += black[i] for v in range(1, MOD): next_i = i + v if next_i >= N: break dp[next_i][v] = dp[i][v] black[next_i] += dp[i][v] black[next_i] %= MOD

しかしこれも、 dpテーブルのサイズが $N \times MOD$ で厳しいです。

気持ちとしては、 $A[i]$ が小さいときにこのdpテーブルをつかいたくて、 $A[i]$ が大きければdpテーブルは使わずに $A[i]$ 刻みで足し算していけます。
こういう $A[i]$ の大小で解き方を変えたいとき、適当にこっちで大小の境界を設定して、 $A[i]$ がそれより大きいときとそうでないときで計算方法を変えてしまっていいです。

詳細は省きますが、その境界はだいたい $\sqrt{N}$ あたりに設定すればいいです。
このあたりは実際に実験してみてよさそうな値を入れてみる、くらいでいいと思います。実際に私はコンテスト本番では境界を $50$ に設定しました( $50$ で実験してみたらうまくいったので)。

コード

MOD = 998244353SQ = 100 # 大小の境界(100くらいがいいかなと思います)N = int(input())A = list(map(int, input().split()))black, white = [0] * N, [0] * Nblack[0] = 1dp = [[0] * (SQ + 1) for _ in range(N)]for i in range(N - 1): if A[i] <= SQ: dp[i][A[i]] += black[i] else: for next_i in range(i + A[i], N, A[i]): black[next_i] += black[i] black[next_i] %= MOD for v in range(1, SQ + 1): next_i = i + v if next_i >= N: break dp[next_i][v] = dp[i][v] black[next_i] += dp[i][v] black[next_i] %= MOD white[i + 1] = (black[i] + white[i]) % MODans = (black[-1] + white[-1]) % MODprint(ans)