ABC241のA,B,C,D,E問題を、Python3でなるべく丁寧に解説していきます。
ただ解けるだけの方法ではなく、次の3つのポイントを満たす解法を解説することを目指しています。
- シンプル：余計なことを考えずに済む
- 実装が楽：ミスやバグが減ってうれしい
- 時間がかからない：パフォが上がって、後の問題に残せる時間が増える
ご質問・ご指摘はコメントかツイッター、マシュマロ、Discordサーバーまでお気軽にどうぞ！
Twitter: u2dayo
マシュマロ： https://marshmallow-qa.com/u2dayo
ほしいものリスト: https://www.amazon.jp/hz/wishlist/ls/2T9IQ8IK9ID19?ref_=wl_share
Discordサーバー（質問や記事の感想・リクエストなどどうぞ！） : https://discord.gg/jZ8pkPRRMT
よかったらLGTMや拡散していただけると喜びます！
目次
ABC241 まとめ
A問題『Digit Machine』
B問題『Pasta』
C問題『Connect 6』
D問題『Sequence Query』
E問題『Putting Candies』
アプリ AtCoderFacts を開発しています
コンテストの統計データを見られるアプリ『AtCoderFacts』を作りました。
現在のところ、次の3つのデータを見ることができます。
- レート別問題正解率
- パフォーマンス目安
- 早解きで上昇するパフォーマンス
今後も機能を追加していく予定です。使ってくれると喜びます。
ABC241 まとめ
全提出人数: 10444人
パフォーマンス
| パフォ | AC | 点数 | 時間 | 順位(Rated内) |
|---|---|---|---|---|
| 200 | AB------ | 300 | 47分 | 6897(6666)位 |
| 400 | AB------ | 300 | 18分 | 5500(5269)位 |
| 600 | AB------ | 300 | 9分 | 4486(4256)位 |
| 800 | ABC----- | 600 | 84分 | 3431(3210)位 |
| 1000 | ABC----- | 600 | 37分 | 2509(2296)位 |
| 1200 | ABCD---- | 1000 | 82分 | 1751(1552)位 |
| 1400 | ABCDE--- | 1500 | 120分 | 1171(986)位 |
| 1600 | ABCDE--- | 1500 | 79分 | 784(605)位 |
| 1800 | ABCDE--- | 1500 | 29分 | 496(347)位 |
| 2000 | ABCDEF-- | 2000 | 92分 | 297(181)位 |
| 2200 | ABCDEF-- | 2000 | 73分 | 178(84)位 |
| 2400 | ABCDE-G- | 2100 | 105分 | 102(38)位 |
色別の正解率
| 色 | 人数 | A | B | C | D | E | F | G | Ex |
|---|---|---|---|---|---|---|---|---|---|
| 灰 | 3616 | 95.2 % | 84.4 % | 17.0 % | 2.6 % | 1.3 % | 0.2 % | 0.0 % | 0.0 % |
| 茶 | 1608 | 98.6 % | 97.6 % | 54.6 % | 14.1 % | 7.4 % | 0.6 % | 0.1 % | 0.0 % |
| 緑 | 1148 | 98.1 % | 97.1 % | 76.5 % | 42.2 % | 32.1 % | 2.1 % | 0.3 % | 0.0 % |
| 水 | 701 | 97.9 % | 97.9 % | 87.6 % | 67.3 % | 71.3 % | 16.6 % | 0.7 % | 0.0 % |
| 青 | 404 | 95.8 % | 95.8 % | 92.1 % | 82.9 % | 92.3 % | 45.8 % | 9.2 % | 1.0 % |
| 黄 | 176 | 86.4 % | 86.4 % | 82.4 % | 75.0 % | 85.2 % | 56.8 % | 27.3 % | 6.2 % |
| 橙 | 38 | 86.8 % | 86.8 % | 86.8 % | 81.6 % | 81.6 % | 68.4 % | 52.6 % | 26.3 % |
| 赤 | 20 | 90.0 % | 95.0 % | 95.0 % | 100.0 % | 95.0 % | 75.0 % | 50.0 % | 35.0 % |
※表示レート、灰に初参加者は含めず
A問題『Digit Machine』
問題ページ：A - Digit Machine
灰コーダー正解率：95.2 %
茶コーダー正解率：98.6 %
緑コーダー正解率：98.1 %
入力
$a_k$ : 画面に数字 $k$ が表示されているとき、ボタンを $1$ 回押すと画面の数字が $a_k$ に変わる
考察
A[A[A[0]]] か、forループで $3$ 回 ans = A[ans] すればいいです。
コード
A = list(map(int, input().split()))
print(A[A[A[0]]])
forループ
A = list(map(int, input().split()))
ans = 0
for _ in range(3):
ans = A[ans]
print(ans)
B問題『Pasta』
問題ページ：B - Pasta
灰コーダー正解率：84.4 %
茶コーダー正解率：97.6 %
緑コーダー正解率：97.1 %
入力
$N$ : $N$ 本の麺のパスタがある
$A_i$ : $i$ 本目の麺の長さは $A_i$
$M$ : $M$ 日間の食事計画がある
$B_i$ : $i$ 日目は 麺の長さがちょうど $B_i$ であるパスタを食べる
考察
各長さの麺が何本残っているか、Counterを使って管理します。$B_i$ が $0$ 本ならば、食事計画を実行できません。$1$ 本以上あるならば、$B_i$ の本数を $1$ 減らします。
Counter同士で引き算が可能なことを利用してもいいです。$B-A$ が空ならばYesです。Counter同士で引き算をすると、$0$ 未満の値は削除されるので、$A-B$ ではうまくいかないことに注意してください。
なお、Python 3.10では <= などの比較演算子を使って部分集合の判定ができるので、B <= A を判定すればいいですが、AtCoderのジャッジサーバー上で動いているPython 3.8 ではできません。
コード
forループ
from collections import Counter
def solve():
N, M = map(int, input().split())
A = Counter(map(int, input().split()))
B = list(map(int, input().split()))
for x in B:
if A[x] == 0:
return False
A[x] -= 1
return True
print("Yes" if solve() else "No")
Counter同士の引き算
from collections import Counter
def solve():
N, M = map(int, input().split())
A = Counter(map(int, input().split()))
B = Counter(map(int, input().split()))
C = B - A # 個数が0になった要素はCounterから削除されます
return not C # Cが空なら False と解釈されるので、not C は Cが空ならTrueを返します
print("Yes" if solve() else "No")
C問題『Connect 6』
問題ページ：C - Connect 6
灰コーダー正解率：17.0 %
茶コーダー正解率：54.6 %
緑コーダー正解率：76.5 %
入力
$N$ : $N\times{N}$ のマス目がある（$N\ge6$）
$S_i$ : $S_i$ の $j$ 文字目は、 $i$ 行目 $j$ 列目のマス目が白か黒かを表す
考察
$N\times{N}$ のマス目上の連続する $6$ マスをすべて調べて、条件を満たすものがあるか判定すればいです。連続するマスは、$4\times{N}\times{N}$ よりも少ないので、十分高速です。
個々の判定は、**『$6$ マスのうち高々 $2$ つ（$2$ つまで）の白マスを黒く塗って、 $6$ マスすべてを黒マスにできるか？』**です。
ある連続する $6$ マスを取り出したとき、 $6$ つすべて黒くできる条件は、$6$ マスに含まれる白マスが $2$ つ以下であることです。
6マスの取り出し方
$N\times{N}$ マスある起点を全探索します。起点から、右、下、左下、右下方向に $6$ マス取り出せばいいです。ただし、マス目からはみ出す場合は判定を行いません。
右、下、左下、右下は $(0,1),(1,0),(1,-1),(1,1)$ 進めばいいです。
コード
PythonではTLEになったので、PyPyで提出してください。（より効率的なコードであれば、Pythonでも通せると思います）
def solve():
def judge(sr, sc, dr, dc):
"""始点(sr, sc) 移動方向(dr, dc)"""
row, col = sr, sc
bl = 0
for _ in range(6):
if not (0 <= row < N and 0 <= col < N): # 6マス取り出す前にマスからはみ出すので判定しない
return False
bl += S[row][col]
row += dr
col += dc
return bl >= 4
pat = [(1, 0), (0, 1), (1, -1), (1, 1)] # 下、右、左下、右下
N = int(input())
S = [[1 if c == "#" else 0 for c in input()] for _ in range(N)] # "#"を1, "."を0に変換します
for row in range(N):
for col in range(N):
for dr, dc in pat:
if judge(row, col, dr, dc):
return True
return False
print("Yes" if solve() else "No")
D問題『Sequence Query』
問題ページ：D - Sequence Query
灰コーダー正解率：2.6 %
茶コーダー正解率：14.1 %
緑コーダー正解率：42.2 %
入力
$Q$ : クエリの数
$\mathrm{query}_i$ : $i$ 番目のクエリは以下の形式
1 x : $A$ に $x$ を追加する。
2 x k : $A$ の $x$ 以下の要素のうち、大きい方から $k$ 番目の値を出力する。
3 x k : $A$ の $x$ 以上の要素のうち、小さい方から $k$ 番目の値を出力する。
- クエリ $2$ と $3$ の $k$ は $5$ 以下
考察
**『順序付き多重集合』**のデータ構造を使えますか？という問題です。
ただし、Pythonの標準ライブラリに『順序付き多重集合』は存在しないので、あらかじめ作成するか、どこかで拾ってくるか、別の手段を使う必要があります。
- C++の標準ライブラリ（STL）にあるmultisetを使う
- Pythonでmultisetを拾ってくる
- クエリ先読み座標圧縮+Binary Indexed Treeで順序付き多重集合シミュレート
解法1 : C++を使う
最も簡単です。upper_boundやlower_boundで二分探索して、あとはイテレータを $k$ 個戻すまたは $k-1$ 個進めるだけです。
コード
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
constexpr ll INF64 = (1LL << 62) - 1;
// 入力高速化 (この問題では300msほど速くなりました）
struct fast_ios {
fast_ios() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
cout << fixed << setprecision(15);
}
} fast_ios;
int main() {
int Q;
cin >> Q;
ll q, x, k;
multiset<ll> S;
for (int i = 0; i < 5; ++i) {
S.insert(-INF64); // 番兵として、-INF64とINF64を5個入れておくと楽です
S.insert(INF64);
}
for (int i = 0; i < Q; ++i) {
cin >> q >> x;
if (q == 1) {
S.insert(x);
}
if (q == 2) {
cin >> k;
auto it = S.upper_bound(x);
for (int j = 0; j < k; ++j) { --it; }
ll ans = *it != -INF64 ? *it : -1;
cout << ans << '\n';
}
if (q == 3) {
cin >> k;
auto it = S.lower_bound(x);
for (int j = 0; j < k - 1; ++j) { ++it; }
ll ans = *it != INF64 ? *it : -1;
cout << ans << '\n';
}
}
}
解法2 : Pythonで順序付き多重集合を使う
@tatyamさんが作ったPythonの順序付き集合/順序付き多重集合を使うといいです。
一般的な順序付き集合は平衡二分探索木などの木構造で実装されることが多いですが、このライブラリは平方分割で実装されています。計算量オーダー的には二分探索木よりも大きいですが、定数倍の関係でこちらのほうが高速に動作するようです。
コード
ライブラリのコードは tatyam-prime/SortedSetを参照してください。
import sys
import math
from bisect import bisect_left, bisect_right, insort
from typing import Generic, Iterable, Iterator, TypeVar, Union, List
readline = sys.stdin.readline # 高速な入力を使います（readlineは改行文字も受け取ることに注意してください）
def main():
Q = int(readline())
S = SortedMultiset()
for _ in range(Q):
query = list(map(int, readline().split()))
q = query[0]
x = query[1]
if q == 1:
S.add(x)
elif q == 2:
k = query[2]
c_le = S.index_right(x) # x以下の要素の数
print(S[c_le - k] if k <= c_le else -1)
else:
k = query[2]
c_lt = S.index(x) # x未満の要素の数
c_ge = len(S) - c_lt # x以上の要素の数
print(S[c_lt + k - 1] if k <= c_ge else -1)
main()
解法3 :クエリ先読み座標圧縮+BIT
順序付き集合や多重集合は、クエリを先読みして座標圧縮したあと、Binary Indexed Tree（Fenwick Tree）上で二分探索をすると同じことができます。
ただし、クエリが先読みできない（オンラインクエリ）場合この方法は使えません。
なお、BIT上の二分探索は普通に二分探索をすると $O((\mathrm{log}N)^2)$ ですが、適切な実装で $O(\mathrm{log}\ N)$ にできます。
コード
import sys
readline = sys.stdin.readline # 高速な入力を使います（readlineは改行文字も受け取ることに注意してください）
def main():
Q = int(readline())
queries = []
A = []
for _ in range(Q):
q = list(map(int, readline().split()))
A.append(q[1])
queries.append(q)
A_sorted = sorted(set(A))
D = {x: i for i, x in enumerate(A_sorted, 1)}
ft = FenwickTree(n=len(D) + 5)
for query in queries:
q = query[0]
xc = D[query[1]]
if q == 1:
ft[xc] += 1
elif q == 2:
k = query[2]
c_le = ft[:xc + 1] # xc以下の要素数
t = c_le - k # 求める値は、全体で小さい方から数えて何番目か
print(A_sorted[ft.upper_bound(t) - 1] if c_le >= k else -1)
else:
k = query[2]
c_lt = ft[:xc] # xc未満の要素数
c_ge = ft[xc:] # xc以上の要素数
t = c_lt + k - 1 # 求める値は、全体で小さい方から数えて何番目か
print(A_sorted[ft.upper_bound(t) - 1] if c_ge >= k else -1)
class FenwickTree:
def __init__(self, n=0, *, array=None):
assert (n > 0 and array is None) or (n == 0 and array)
if array:
n = len(array)
self.__array = array[:] # get用
_container = array[:]
for i in range(n):
j = i | (i + 1)
if j < n:
_container[j] += _container[i]
self.__size = len(array) + 1
self.__container = [0] + _container[:]
self.__depth = n.bit_length()
else:
self.__array = [0] * n
self.__size = n + 1
self.__container = [0] * (n + 1)
self.__depth = n.bit_length()
def add(self, i, x):
"""a[i]にxを加算"""
self.__array[i] += x
i += 1
while i < self.__size:
self.__container[i] += x
i += i & (-i)
def sum(self, r):
"""[0, r) の総和"""
s = 0
while r > 0:
s += self.__container[r]
r -= r & (-r)
return s
def sum_range(self, l, r):
"""[l, r) の総和"""
return self.sum(r) - self.sum(l)
def upper_bound(self, s):
"""[0, r) の総和 <= s となる最大のrを求める"""
w, r = 0, 0
for i in reversed(range(self.__depth)):
k = r + (1 << i)
if k < self.__size and w + self.__container[k] <= s:
w += self.__container[k]
r += 1 << i
return r
def set(self, i, x):
"""a[i]をxに変更"""
self.add(i, x - self.__array[i])
def get(self, i):
"""a[i]を返す"""
return self.__array[i]
def __getitem__(self, key):
if isinstance(key, slice):
start, stop = key.start, key.stop
if start is None: start = 0
if stop is None: stop = self.__size - 1
if start == 0: return self.sum(stop)
return self.sum_range(start, stop)
else:
return self.get(key)
def __setitem__(self, key, value):
self.set(key, value)
if __name__ == '__main__':
main()
E問題『Putting Candies』
問題ページ：E - Putting Candies
灰コーダー正解率：1.3 %
茶コーダー正解率：7.4 %
緑コーダー正解率：32.1 %
入力
$N$ : 長さ $N$ の 正整数列が与えられる
$K$ : 操作を $K$ 回行ったあとのアメの個数を答える
$A_i$ : $i$ 番目の皿には $A_i$ 個のアメが乗っている
考察
**『皿の中のアメの数 $X$ を $N$ で割った値』のことを『店』**と呼ぶことにします。
店 $i$ にいるとき、次の操作では アメを $A_i$ 個得て、店 $(X+A_i)\ \mathrm{mod}\ N$ に移動します。開始時点では、店 $0$ にいます。次に行く店は、$0+A_0$ を $N$ で割った場所です。
以下の重要な性質があります。
- 店 $i$ にいるとき、$X$ を $N$ で割った値は常に $i$
この性質より、さらに重要な性質が得られます。
- 店 $i$ から次に行く店 $j$ は、操作の回数や $X$ そのものにはよらず、常に $(i+A_i)\ \mathrm{mod}\ N$ である
解法1: ダブリング
ダブリングで解くことができます。
- 店 $i$ から $2^k$ 回操作を行ったとき、どの店 $j$ に行くか？
- 店 $i$ から $2^k$ 回操作を行ったとき、アメの総数 $X$ はいくつ増えるか？
は、$2^k-1$ 回操作を行ったとき……の情報より求めることができます。この $2$ つの情報を、あらかじめ作っておくか、計算過程で都度更新していけばいいです。
そして、$K$ を $2$ 進法表記したときの $i$ 桁目が $1$ であれば、現在の店とアメの総数を更新すればいいです。（繰り返し二乗法の要領です）
コード
どちらのコードもPythonではTLEになったので、PyPyで出してください。PyPyの場合実行時間に大きな差は見られませんでした。
先にテーブルを作る
def main():
N, K = map(int, input().split())
A = list(map(int, input().split()))
next_dbl = [[(i + x) % N for i, x in enumerate(A)]]
sum_dbl = [A[:]]
s = 2
while s <= K:
next_curr = [0] * N
sum_curr = [0] * N
next_prev = next_dbl[-1]
sum_prev = sum_dbl[-1]
for i in range(N):
next_curr[i] = next_prev[next_prev[i]]
sum_curr[i] = sum_prev[i] + sum_prev[next_prev[i]]
next_dbl.append(next_curr)
sum_dbl.append(sum_curr)
s *= 2
ans = 0
curr = 0
for i in range(len(sum_dbl)):
if K >> i & 1:
ans += sum_dbl[i][curr]
curr = next_dbl[i][curr]
print(ans)
if __name__ == '__main__':
main()
都度更新
PythonではTLEですが、先ほどのコードよりもこちらのコードのほうが高速でした。
def main():
N, K = map(int, input().split())
A = list(map(int, input().split()))
next_dbl = [(i + x) % N for i, x in enumerate(A)]
sum_dbl = A[:]
ans = 0
curr = 0
while K > 0:
if K & 1:
ans += sum_dbl[curr]
curr = next_dbl[curr]
next_new = [0] * N
sum_new = [0] * N
for i in range(N):
next_new[i] = next_dbl[next_dbl[i]]
sum_new[i] = sum_dbl[i] + sum_dbl[next_dbl[i]]
next_dbl = next_new
sum_dbl = sum_new
K >>= 1
print(ans)
if __name__ == '__main__':
main()
解法2: ループすることを利用する
訪れる店は一定で、周期性があるということは、いつかは同じところをグルグル回り続けるサイクルに入ります。$\rho$ の字のようなイメージです。
- サイクルに入る前に訪れる店を順番に
- サイクルで訪れる店を順番に
- サイクル $1$ 周でいくつアメを得られるか（サイクル中で訪れる店の $A_i$ の合計）
- サイクルを何周するか、$1$ 周に満たない操作は何回か（割り算の商と余りです）
がわかれば、$K$ が巨大でもこの問題を解くことができます。
自分で『訪れた店を記録して、$2$ 回訪れたらサイクルに入っている』のように実装してもいいですが、フロイドの循環検出法という空間計算量 $O(1)$ のアルゴリズムがあるので、それをライブラリにしておくと楽です。（※なお以下のコードはFloyd's cycle-finding algorithmの改良版であるBrent's algorithm です）
def main():
def solve():
def f(x):
return (x + A[x % N]) % N
def calc_sum(L):
ret = 0
for t in L:
ret += A[t]
return ret
def naive():
ans = 0
rem = K
while rem > 0:
ans += A[ans % N]
rem -= 1
return ans
N, K = map(int, input().split())
A = list(map(int, input().split()))
# サイクルに入る前に終わるパターンの場合分けが面倒なので、Kが小さい場合は愚直解を返す
if K <= N + 10: return naive()
ans = 0
rem = K
bfc, cyc = cycle_detection_lists(f, 0) # サイクルに入る前に訪れる店・サイクル中の店のリストを返す
# サイクルに入る前の計算
rem -= len(bfc)
ans += calc_sum(bfc)
# サイクルに入った後の計算
q, r = divmod(rem, len(cyc)) # サイクル全体q回と、サイクルの最初から途中までr個
sum_cyc = calc_sum(cyc) # サイクル1周でアメをいくつ得られるか
ans += sum_cyc * q
ans += calc_sum(cyc[:r])
return ans
print(solve())
def cycle_detection(f, x0):
"""
数列 A_{i+1} = f(A_i) のループの始点と周期を求める
f: 関数 f(x)
x0: A_0
mu: ループの始点 A_{mu}
lam: ループの周期 A_{mu} = A_{mu + k * lam}
"""
power = lam = 1
first = x0
second = f(x0)
while first != second:
if power == lam:
first = second
power *= 2
lam = 0
second = f(second)
lam += 1
first = second = x0
for i in range(lam):
second = f(second)
mu = 0
while first != second:
first = f(first)
second = f(second)
mu += 1
return mu, lam
def cycle_detection_lists(f, x0):
"""サイクルに入る前・サイクル中の数列のそのものを返す"""
mu, lam = cycle_detection(f, x0)
before_cycle = [0] * mu
if mu > 0:
before_cycle[0] = x0
for i in range(mu - 1):
before_cycle[i + 1] = f(before_cycle[i])
cycle = [0] * lam
curr = x0 if mu == 0 else f(before_cycle[-1])
cycle[0] = curr
for i in range(lam - 1):
cycle[i + 1] = f(cycle[i])
return before_cycle, cycle
if __name__ == '__main__':
main()