AtCoder Beginners Contest 343 (ABC343) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。
A - Wrong Answer
問題
$0$ 以上 $9$ 以下の整数 $A, B$ が与えられます。
$0$ 以上 $9$ 以下の整数のうち、$A+B$ と等しくない数をひとつ出力してください。
考察
for文を用いて、$x=0,1,2, \cdots ,9$ と探索します。
$x \neq A+B$ となる $x$ があれば、その $x$ を出力し、コードを終了します。
コード
A, B = map(int, input().split())
for x in range(10):
if A + B != x:
print(x)
break
B - Adjacency Matrix
問題
$N$ 頂点の単純無向グラフ $G$ があります。
また、$N$ 行 $N$ 列の行列 $A$ が与えられます。$2$ 頂点 $i,j$ 間に辺があるとき $A_{i,j}=1$ 、辺がないとき $A_{i,j}=0$ です。
$i=1,2, \cdots ,N$ について、頂点 $i$ と辺で直接結ばれた頂点番号を昇順に出力してください。
考察
$N, A$ を入力で受け取ったあと、$i=1,2, \cdots ,N$ についてそれぞれ独立に考えます。
頂点 $i$ と辺で直接結ばれた頂点は、行列 $A$ の $i$ 行目を見ればわかります。
行列 $A$ の $i$ 行目について、 $j$ 列目( $j=1,2,\cdots ,N$ )の値が $1$ のとき、頂点 $i,j$ は辺で直接結ばれています。
コード
N = int(input())
A = [list(map(int, input().split())) for _ in range(N)]
for i in range(N):
ans_lst = []
for j in range(N):
if A[i][j] == 1:
ans_lst.append(j + 1)
print(*ans_lst)
C - 343
問題
正整数 $N$ が与えられます。
$N$ 以下で最大の回文立方数を出力してください。
ただし $v$ が回文立方数であるとは、以下の $2$ 条件がどちらも満たされてることを指します。
- $v = t^3$ となる正整数 $t$ が存在する
- $v$ を文字列として見たとき、その文字列が回文になっている
考察
制約が $1 \leq N \leq 10^{18}$ なので、for文をつかって $v=1,2,\cdots , 10^{18}$ について回文立方数かどうかを確認するとTLEしてしまいます。
$N$ 以下で最大の回文立方数が $v$ のとき、 $v = t^3$ となる正整数 $t$ が存在しています。 $v \leq N \leq 10^{18}$ なので、$t^3 \leq 10^{18}$ が成り立ちます。 $t$ について解くと $t \leq 10^6$ となります。
for文をつかって $v=1,2, \cdots 10^{18}$ について回文立方数かどうかを確認する代わりに、 for文をつかって $t=1,2,\cdots 10^6$ について、$v = t^3$ の値が回文になっているかどうかを確認します。
つまり、問題を次のように言い換えて解きます。
正整数 $N$ が与えられます。
以下の $2$ 条件をともに満たす最大の $t$ の値を求め、$t^3$ を出力してください。
- $v = t^3$ としたとき、 $v \leq N$ である
- $v$ を文字列として見たとき、回文になっている
回文かどうかの判定は、メインのコードとは別で関数化しておくとコードがスッキリしてオススメです。
コード
# 文字列として見たとき、回文になっていればTrueを、そうでなければFalseを返す。
def is_kaibun(v):
str_v = str(v)
if str_v == str_v[::-1]:
return True
else:
return False
N = int(input())
ans = 1
for t in range(1, 1_000_001):
v = t ** 3
if v > N:
break
if is_kaibun(v):
ans = v
print(ans)
D - Diversity of Scores
問題
$N$ 人の参加者がいて、最初は全員が持ち点 $0$ です。
$i=1,2, \cdots ,T$ について、$i$ 秒後に選手 $A_i$ が持ち点を $B_i$ だけ増やすことが分かっています。
$i=1,2,\cdots ,T$ について、$i+0.5$ 秒後における $N$ 人の持ち点の種類数を出力してください。
考察
キーが持ち点・バリューがその持ち点である参加者の数、になるような辞書 $dic$ を用意します。
また、参加者 $x$ の持ち点が $y$ のとき $P[x] = y$ になるようなリスト $P$ を用意します。
辞書 $dic$ とリスト $P$ を管理しながら、 $T$ 個のクエリを順に処理します。
持ち点の種類数は、辞書 $dic$ のキーの種類数と等しいです。これは len(dic) で求められます。
コード
N, T = map(int, input().split())
dic = dict()
dic[0] = N
P = [0] * N
for _ in range(T):
a, b = map(int, input().split())
a -= 1
dic[P[a]] -= 1
if dic[P[a]] == 0:
del dic[P[a]]
P[a] += b
if P[a] not in dic:
dic[P[a]] = 0
dic[P[a]] += 1
print(len(dic))
E - 7x7x7
問題
$1$ 辺の長さ $7$ の立方体が $3$ つあります。
$i=1,2,3$ について、 $i$ 個の立方体に覆われた体積が $V_i$ になるように立方体を配置したいです。
ただし、各辺は $x,y,z$ 軸いずれかと平行で、立方体の角の座標 $(a,b,c)$ について $a,b,c$ はすべて整数になるように配置するものとします。
考察1 (配置のパターン数)
以下、$a \leq x \leq a+7, \enspace b \leq y \leq b+7, \enspace c \leq z \leq c+7$ の領域に立方体を配置することを、$(a,b,c)$ に配置すると表現します。
とりあえず $1$ つ目の立方体はどこに配置してもいいので、 $(0,0,0)$ に配置することにします。
$2$ つ目の立方体は、$1$ つ目の立方体との相対的な位置だけが重要なので、たとえば $(-1,0,0)$ に配置するのと $(1,0,0)$ に配置するのは同じです。
なので、$2$ つ目の立方体の座標 $(x_2, y_2, z_2)$ について、 $0 \leq x_2, y_2, z_2 \leq 7$ の範囲にしか置かないとしても大丈夫です。
$1,2$ つ目の立方体で覆われた領域は広く見積もっても $0 \leq x,y,z \leq 14$ の範囲です。なので、$3$ つ目の立方体の座標 $(x_3,y_3, z_3)$ について、 $-7 \leq x,y,z \leq 14$ の範囲にしか置かないとしても大丈夫です。
$2$ つ目の立方体の配置は $8^3$ 通り、$3$ つ目の立方体の配置は $22^3$ 通りです。全体として配置は $8^3 \times 22^3$ 通り考えればよく、計算すると $8^3 \times 22^3 = 5451776 \lt 5.5 \times 10^6$ です。
考察2 (重なっている領域の体積を求める)
$2$ つの立方体の座標が $(x_1, y_1, z_1), \enspace (x_2, y_2, z_2)$ のときを考えます。
この $2$ つが重なっている領域は直方体になります。この直方体の各辺の長さを求めることで、重なっている領域の体積を求めます。
まず $x$ 軸方向のみを考えます。重なっている領域の $x$ について、 $x_1 \leq x \leq x_1 + 7, \enspace x_2 \leq x \leq x_2 + 7$ を同時に満たしています。
まとめて書くと、 $\max (x_1, x_2) \leq x \leq \min (x_1, x_2) + 7$ です。
$y, z$ 軸方向に関しても同様です。
また、$3$ つの立方体が重なっている領域も同じように考えることができます。
$x$ 軸方向に関しては、 $\max (x_1, x_2, x_3) \leq x \leq \min (x_1, x_2, x_3) + 7$ です。
$y, z$ 軸方向に関しても同様です。
立方体の座標が分かっているとき、この計算は $O(1)$ でできます。
考察3 (解く)
考察2での計算は関数化しておきます。
def common_length2(x1, x2):
left = max(x1, x2)
right = min(x1, x2) + 7
if left <= right:
return right - left
else:
return 0
def common_length3(x1, x2, x3):
left = max(x1, x2, x3)
right = min(x1, x2, x3) + 7
if left <= right:
return right - left
else:
return 0
def calc_volume3(a1, b1, c1, a2, b2, c2, a3, b3, c3):
dx = common_length3(a1, a2, a3)
dy = common_length3(b1, b2, b3)
dz = common_length3(c1, c2, c3)
return dx * dy * dz
def calc_volume2(a1, b1, c1, a2, b2, c2):
dx = common_length2(a1, a2)
dy = common_length2(b1, b2)
dz = common_length2(c1, c2)
return dx * dy * dz
立方体がちょうど $2$ つ重なっている領域の体積について、包除原理と同じ要領で次のように計算できます。
volume3 = calc_volume3(a1, b1, c1, a2, b2, c2, a3, b3, c3)
volume2_1 = calc_volume2(a1, b1, c1, a2, b2, c2)
volume2_2 = calc_volume2(a1, b1, c1, a3, b3, c3)
volume2_3 = calc_volume2(a2, b2, c2, a3, b3, c3)
volume2 = volume2_1 + volume2_2 + volume2_3 - volume3 * 3
ちょうど $1$ つだけの立方体に覆われた領域の体積は、同じ考え方で計算できます。
volume1 = (7 * 7 * 7 * 3) - volume2 * 2 - volume3 * 3
考察1で配置のパターン数が $5.5 \times 10^6$ 個程度、考察2での計算は $O(1)$ で、計算量は問題ありません。
これをまとめると、下のコードになります。
(TLEしてしまいます、続きはおまけを見てね。)
コード
"""注意
TLEコードです。
"""
def common_length2(x1, x2):
left = max(x1, x2)
right = min(x1, x2) + 7
if left <= right:
return right - left
else:
return 0
def common_length3(x1, x2, x3):
left = max(x1, x2, x3)
right = min(x1, x2, x3) + 7
if left <= right:
return right - left
else:
return 0
def calc_volume3(a1, b1, c1, a2, b2, c2, a3, b3, c3):
dx = common_length3(a1, a2, a3)
dy = common_length3(b1, b2, b3)
dz = common_length3(c1, c2, c3)
return dx * dy * dz
def calc_volume2(a1, b1, c1, a2, b2, c2):
dx = common_length2(a1, a2)
dy = common_length2(b1, b2)
dz = common_length2(c1, c2)
return dx * dy * dz
def calc_volume(a1, b1, c1, a2, b2, c2, a3, b3, c3):
volume3 = calc_volume3(a1, b1, c1, a2, b2, c2, a3, b3, c3)
volume2_1 = calc_volume2(a1, b1, c1, a2, b2, c2)
volume2_2 = calc_volume2(a1, b1, c1, a3, b3, c3)
volume2_3 = calc_volume2(a2, b2, c2, a3, b3, c3)
volume2 = volume2_1 + volume2_2 + volume2_3 - volume3 * 3
volume1 = (7 * 7 * 7 * 3) - volume2 * 2 - volume3 * 3
return volume1, volume2, volume3
v1, v2, v3 = map(int, input().split())
a1, b1, c1 = 0, 0, 0
for a2 in range(8):
for b2 in range(8):
for c2 in range(8):
for a3 in range(-7, 15):
for b3 in range(-7, 15):
for c3 in range(-7, 15):
if calc_volume(a1, b1, c1, a2, b2, c2, a3, b3, c3) == (
v1, v2, v3):
print("Yes")
print(a1, b1, c1, a2, b2, c2, a3, b3, c3)
exit()
print("No")
おまけ
計算量自体は問題ないですが、for文でネストがかなり深くなっているのと体積の計算が $O(1)$ の割にはそこそこ計算しているので、上のコードだとTLEしてしまいます。
別の高速な言語で書き換えましょう。
chatGPTを利用して、次のようにAIに質問します。
「次のPythonコードをRustに書き換えてください。入力の受け取りはproconio::inputを使用してください。(以下、先ほどのTLEコード)」
単純で簡潔なコードなら、ちゃんと書きかえてくれます。
F - Second Largest Query
問題
数列 $A=(A_1, A_2, \cdots A_N)$ が与えられます。
$Q$ 個のクエリを順に処理してください。クエリは以下の $2$ 種類です。
-
1 p x: $A_p$ の値を $x$ に変更する -
2 l r: $A_l, A_{l+1}, \cdots A_r$ の中で $2$ 番目に大きい値の個数を出力する
考察
セグメントツリーに次のタプルをのせます。
$$(1番目に大きい数 x_1,\enspace x_1の個数,\enspace 2番目に大きい数x_2,\enspace x_2の個数)$$
要素数 $4$ のタプルをのせるので1つ1つの計算がそこそこ重くなってしまいます。適当にコードを書くとすぐTLEするので、なるべく簡単なコードを書くようにしましょう。
下のコードでは、関数 $op$ の中で気合い入れて場合分けしておくことで、Python側の計算回数が少なくなるようにしています。
コード
from atcoder.segtree import SegTree
def op(ele1, ele2):
v11, c11, v12, c12 = ele1
v21, c21, v22, c22 = ele2
if v11 > v21:
if v12 > v21:
return v11, c11, v12, c12
elif v12 < v21:
return v11, c11, v21, c21
else:
return v11, c11, v12, c12 + c21
elif v11 < v21:
if v11 > v22:
return v21, c21, v11, c11
elif v11 < v22:
return v21, c21, v22, c22
else:
return v21, c21, v11, c11 + c22
else:
if v12 > v22:
return v11, c11 + c21, v12, c12
elif v12 < v22:
return v11, c11 + c21, v22, c22
else:
return v11, c11 + c21, v12, c12 + c22
e = (-1, 0, -2, 0) # opの単位元
N, Q = map(int, input().split())
A = list(map(int, input().split()))
seg = SegTree(op, e, [(a, 1, -1, 0) for a in A])
for _ in range(Q):
q, a, b = map(int, input().split())
if q == 1:
p, x = a - 1, b
seg.set(p, (x, 1, -1, 0))
else:
l, r = a - 1, b
print(seg.prod(l, r)[3])
おまけ
たまたま上のコードは1400msくらいでACできていますが、少しでもより重くなるようなコードを書いたり問題設定が厳しかったりするとTLEしてしまいます。
セグメントツリーで重めの計算をするときや遅延セグメントツリーをつかうとき、C++やRustなどの高速な言語が書けるならそっちで書いた方がいいでしょう。