ABC335(A~D)をPythonで解く

A: 2023

スライスで文字列$S$の最後の一文字を取り除き、4を挿入します。

コード

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

B: Tetrahedral Number

$N\leq21$と十分に小さいので、3重ループを回して全探索が可能です。
辞書順に小さいほうから全て出力するのは $x$, $y$, $z$の順番に探索すれば十分ですが、念のため最後にソートすると安心です。

$N$を含むようループのrangeに注意してください。

コード

N = int(input())
ans = []
for x in range(N+1):
    for y in range(N+1):
        for z in range(N+1):
            if x+y+z <= N:
                ans.append([x,y,z])

ans = sorted(ans)
for i in ans:
    print(*i)

C: Loong Tracking

頭以外の全てのパーツは頭が通った軌跡を辿るので、頭の軌跡を調べれば良いです。

龍の初期位置のリストに頭が動いていく座標を追加していくのですが、リストの前に追加すると挿入で$O(N)$かかります。なので後ろにappendしていきましょう。

最後尾が通った後の座標は削除するとまた$O(N)$かかるのでそのまま放置しておきます。

従ってこの問題は以下のように解くことができます。

1. 初期位置を頭が後ろに来るようにリストに持つ
2. 頭が動いた先の座標をリストにappend
3. 座標を訊かれたら後ろから数えて出力

コード

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

# 龍の初期位置
point = [(i, 0) for i in range(N, 0, -1)]

for _ in range(Q):
    t, c = map(str, input().split())

    # 頭の向かう先の座標
    if t == "1":
        x,y = point[-1]
        if c == "U":
            y += 1
        elif c == "D":
            y -= 1
        elif c == "L":
            x -= 1
        else:
            x += 1
        point.append((x, y))
    
    # 後ろから数えて出力
    else:
        print(*point[-int(c)])

また、座標の追加と削除はcollectionsのdequeにてそれぞれ$O(1)$で可能ですが、これを使うと座標出力（dequeへのランダムアクセス）の際に$O(N)$かかってしまいTLEします。

リストのappendとランダムアクセスは共に$O(1)$であり十分高速です。

代替法として、ごりちゃんさんのブログに$O(1)$でランダムアクセスが可能なdequeの実装があります。

ごりちゃんさんdequeを使った場合のコード

class Deque:
    def __init__(self, src_arr=[], max_size=300000):
        self.N = max(max_size, len(src_arr)) + 1
        self.buf = list(src_arr) + [None] * (self.N - len(src_arr))
        self.head = 0
        self.tail = len(src_arr)
    def __index(self, i):
        l = len(self)
        if not -l <= i < l: raise IndexError('index out of range: ' + str(i))
        if i < 0:
            i += l
        return (self.head + i) % self.N
    def __extend(self):
        ex = self.N - 1
        self.buf[self.tail+1 : self.tail+1] = [None] * ex
        self.N = len(self.buf)
        if self.head > 0:
            self.head += ex
    def is_full(self):
        return len(self) >= self.N - 1
    def is_empty(self):
        return len(self) == 0
    def append(self, x):
        if self.is_full(): self.__extend()
        self.buf[self.tail] = x
        self.tail += 1
        self.tail %= self.N
    def appendleft(self, x):
        if self.is_full(): self.__extend()
        self.buf[(self.head - 1) % self.N] = x
        self.head -= 1
        self.head %= self.N
    def pop(self):
        if self.is_empty(): raise IndexError('pop() when buffer is empty')
        ret = self.buf[(self.tail - 1) % self.N]
        self.tail -= 1
        self.tail %= self.N
        return ret
    def popleft(self):
        if self.is_empty(): raise IndexError('popleft() when buffer is empty')
        ret = self.buf[self.head]
        self.head += 1
        self.head %= self.N
        return ret
    def __len__(self):
        return (self.tail - self.head) % self.N
    def __getitem__(self, key):
        return self.buf[self.__index(key)]
    def __setitem__(self, key, value):
        self.buf[self.__index(key)] = value
    def __str__(self):
        return 'Deque({0})'.format(str(list(self)))


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

point = Deque([(i, 0) for i in range(N, 0, -1)])

for _ in range(Q):
    t, c = map(str, input().split())

    if t == "1":
        x,y = point[-1]
        if c == "U":
            y += 1
        elif c == "D":
            y -= 1
        elif c == "L":
            x -= 1
        else:
            x += 1
        point.append((x, y))

        # このpopleftは無くてもいいです
        # 最後尾通過済みの座標を削除
        point.popleft()
    
    else:
        print(*point[-int(c)])

D: Loong and Takahashi

問題文の条件はいろいろ書いてありますが、よく解らないので出力例を見るとうずまき状に数字を配置すればよさそうということがわかります。

ではどううずまきが実装できるか考えてみます。

まずは$N$x$N$の行列を準備しておきます。初めは全てのマスを0で初期化しておくと良さそうです。

左上から始めるとして、

1. 突き当りまで右に進む
2. 突き当りまで下に進む
3. 突き当りまで左に進む
4. 突き当りまで上に進む
5. ...

を繰り返せばよいことがわかります。

ここで上下左右移動の座標変動（dx, dy）を考えてみると、

右: (1, 0)
下: (0, 1)
左: (-1, 0)
上: (0, -1)

となります。

(よく混乱しやすいことですが、xy平面上での座標移動と行列内での移動は上下の向きが逆になります。xy平面上での下移動は$y$座標を減らす (-1) ことを意味しますが、行列内での下移動は次のリスト(row)に移る (+1) を意味するからです。)

この変化を表すために、

dx, dy = -dy, dx

で 右→下→左→上 の移動を実装することができます。

matrix = [[0] * N for _ in range(N)]

# 左上から始める
x, y = 0, 0

# 最初は右方向への移動
dx, dy = 1, 0


for i in range(1, n * n + 1):

    matrix[y][x] = i

    # not(次のマスが存在する、かつまだ埋められていないとき)
    # つまり「次のマスが存在しなかったりもう埋められていたとき」
    if not (0 <= x + dx < N and 0 <= y + dy < N and matrix[y+dy][dx+x] == 0):
        # 方向を変えます
        dx, dy = -dy, dx
    
    x += dx
    y += dy

真ん中はTを置くので、座標を計算して置き換えます。

center = N // 2
matrix[center][center] = 'T'

あとは出力してやれば終わりです。

コード

def create_spiral_matrix(N):
    matrix = [[0] * N for _ in range(N)]

    # 左上から始める
    x, y = 0, 0

    # 最初は右方向への移動
    dx, dy = 1, 0

    for i in range(1, n * n + 1):

        matrix[y][x] = i

        # not(次のマスが存在する、かつまだ埋められていないとき)
        # つまり「次のマスが存在しなかったりもう埋められていたとき」
        if not (0 <= x + dx < N and 0 <= y + dy < N and matrix[y+dy][dx+x] == 0):
            # 方向を変えます
            dx, dy = -dy, dx
        
        x += dx
        y += dy

    center = N // 2
    matrix[center][center] = 'T'

    return matrix

matrix = create_spiral_matrix(int(input()))
for row in matrix:
    print(*row)

E: Non-Decreasing Colorful Path

TODO