ABC228のA,B,C,D問題を、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や拡散していただけると喜びます！
目次
ABC228 まとめ
A問題『On and Off』
B問題『Takahashi's Secret』
C問題『Final Day』
D問題『Linear Probing』
アプリ AtCoderFacts を開発しています
コンテストの統計データを見られるアプリ『AtCoderFacts』を作りました。
現在のところ、次の3つのデータを見ることができます。
- レート別問題正解率
- パフォーマンス目安
- 早解きで上昇するパフォーマンス
今後も機能を追加していく予定です。使ってくれると喜びます。
ABC228 まとめ
全提出人数: 6298人
パフォーマンス
| パフォ | AC | 点数 | 時間 | 順位(Rated内) |
|---|---|---|---|---|
| 200 | AB------ | 300 | 45分 | 4794(4502)位 |
| 400 | ABC----- | 600 | 91分 | 3979(3690)位 |
| 600 | ABC----- | 600 | 60分 | 3317(3028)位 |
| 800 | ABC----- | 600 | 38分 | 2635(2348)位 |
| 1000 | ABC----- | 600 | 15分 | 2010(1724)位 |
| 1200 | ABCD---- | 1000 | 78分 | 1485(1205)位 |
| 1400 | ABCD---- | 1000 | 44分 | 1070(799)位 |
| 1600 | ABCDE--- | 1500 | 149分 | 746(501)位 |
| 1800 | ABCDE--- | 1500 | 94分 | 516(290)位 |
| 2000 | ABCDE--- | 1500 | 69分 | 353(153)位 |
| 2200 | ABCDE--- | 1500 | 37分 | 246(74)位 |
| 2400 | ABCDEF-- | 2000 | 96分 | 156(34)位 |
色別の正解率
| 色 | 人数 | A | B | C | D | E | F | G | H |
|---|---|---|---|---|---|---|---|---|---|
| 灰 | 2519 | 81.3 % | 70.0 % | 46.6 % | 3.9 % | 0.5 % | 0.0 % | 0.0 % | 0.0 % |
| 茶 | 1145 | 94.8 % | 97.3 % | 89.0 % | 20.3 % | 3.6 % | 0.1 % | 0.0 % | 0.0 % |
| 緑 | 963 | 97.3 % | 98.5 % | 93.7 % | 49.7 % | 11.2 % | 0.5 % | 0.0 % | 0.0 % |
| 水 | 573 | 98.1 % | 99.3 % | 96.5 % | 80.5 % | 34.0 % | 3.0 % | 0.0 % | 0.3 % |
| 青 | 337 | 99.1 % | 99.7 % | 99.7 % | 92.3 % | 65.9 % | 16.6 % | 1.2 % | 0.9 % |
| 黄 | 231 | 94.8 % | 94.8 % | 93.1 % | 89.2 % | 81.0 % | 48.0 % | 6.1 % | 7.4 % |
| 橙 | 43 | 95.3 % | 95.3 % | 95.3 % | 95.3 % | 90.7 % | 81.4 % | 44.2 % | 41.9 % |
| 赤 | 23 | 100.0 % | 100.0 % | 100.0 % | 95.7 % | 100.0 % | 91.3 % | 95.7 % | 69.6 % |
※表示レート、灰に初参加者は含めず
A問題『On and Off』
問題ページ：A - On and Off
灰コーダー正解率：81.3 %
茶コーダー正解率：94.8 %
緑コーダー正解率：97.3 %
史上最高に難しいA問題でした。私も5分と1WAを使いました。
入力
$S$ : $S$ 時 $0$ 分に電気がつく
$T$ : $T$ 時 $0$ 分に電気が消える
$X$: $X$ 時 $30$ 分に電気がついているか判定する
※電気がついている間に日付が変わる場合がある（$S<T$ とは限らない）
考察
日付をまたがない場合（S < T)
日付をまたがない場合（$S<T$）を考えます。$S$ 時 $30$ 分には電気がついています。$T-1$ 時 $30$ 分までは電気がついていますが、$T$ 時 $30$ 分には既に電気が消えています。
よって、電気がついている条件は $S\le{X}\lt{T}$ です。$T$ は含みません。
日付をまたぐ場合(T>S)
入力例 $2$ にあるように、電気をついている間に日付をまたぐ場合があります。
入力例 $2$ の $S=20, T=7, X=12$ に、日付をまたぐ場合の丁寧な説明があります。
部屋の電気がついているのは $0$ 時 $0$ 分から $7$ 時 $0$ 分までの間と、$20$ 時 $0$ 分から（次の日の）$0$ 時 $0$ 分までの間です。
この説明通りに条件を書くと、$0\le{X}\lt{7}$ または $20\le{X}\lt{24}$ のどちらかであれば電気がついていることになります。この例では $X=12$ ですから、どちらも満たさないため No になります。
電気がついている条件を $S,T$ を使って一般化すると、$0\le{X}\lt{T}$ または $S\le{X}\lt{24}$ になります。
なお、$0\le{X}\le{23}$ という制約があるため、$X\lt{T}$ または $S\le{X}$ と簡略化できます。
他の考え方：電気がついていないか判定して、反転する
『電気がついている間に日付をまたぐ』ならば『電気が消えている間は日付をまたがない』です。そこで、$X$ 時 $30$ 分に電気がついていないかを判定することでも日付をまたぐ場合を解くことができます。
$X$ 時 $30$ 分に電気がついていない条件は、$T\le{X}\lt{S}$ です。これを満たすならNoで、満たさないならYesです。
コード
コード1
S, T, X = map(int, input().split())
if S < T: # 日付が変わらない場合
print("Yes" if S <= X < T else "No") # T時30分に電気は消えていますから、Tは含みません
else: # 日付が変わる場合
print("Yes" if X < T or S <= X else "No")
コード2
S, T, X = map(int, input().split())
if S < T: # 日付が変わらない場合
print("Yes" if S <= X < T else "No") # T時30分に電気は消えていますから、Tは含みません
else: # 日付が変わる場合
print("Yes" if not (T <= X < S) else "No") # X時30分に電気が消えているか判定しています
B問題『Takahashi's Secret』
問題ページ：B - Takahashi's Secret
灰コーダー正解率：70.0 %
茶コーダー正解率：97.3 %
緑コーダー正解率：98.5 %
入力
$N$ : 友達の人数
$X$ : はじめに高橋君の秘密を知った友達の番号が $X$
$A_i$ : 友達 $i$ が秘密を知ったとき、友達 $A_i$ に秘密を教える
考察
伝言ゲームのようなに秘密が広まっていくイメージの問題です。
はじめ、『秘密を知らない友達の総数』は、友達 $X$ を除いた $N-1$ 人います。誰かが『秘密を知らない友達』に秘密を伝えると、『秘密を知らない友達の総数』は $1$ 人減ります。（そのままのことを言っていますね）
いつかは必ず『秘密を知っている友達の総数』は $0$ 人になる（全員が秘密を知る）ため、伝言ゲームは必ず終わります。もちろん、全員が秘密を知る前に、秘密を知っている友達に行き着いて伝言ゲームが終了する場合もあります。
ですから、問題文に書いてある通りに伝言ゲームが終わるまでシミュレーションをして、最終的に秘密を知った友達の数をカウントすればいいです。
コード
N, X = map(int, input().split())
A = [-1] + list(map(int, input().split())) # A[1]が友達1の情報を表すようにしたいので、A[0]にダミーを入れてずらします
flag = [False] * (N + 1) # flag[i] が True なら、友達iは秘密を知っています
flag[X] = True
now = X # はじめに秘密を知ったのは友達Xです
while not flag[A[now]]: # 友達A_iがまだ秘密を知らなければ、A_iに伝えます（知っていればそこで終了です）
now = A[now]
flag[now] = True
print(flag.count(True))
C問題『Final Day』
問題ページ：C - Final Day
灰コーダー正解率：46.6 %
茶コーダー正解率：89.0 %
緑コーダー正解率：93.7 %
入力
$N$ : 生徒の数
$K$ : 各生徒について、$4$ 日目の試験終了後に上位 $K$ 位以内に入っていることがあり得るか判定する
$P_{i,j}$ 生徒 $i$ の $j$ 日目の試験の点数（それぞれ $0$ ～ $300$ 点）
考察
$3$ 日目までの結果は既に確定していますから、変えることはできません。
ある生徒 $x$ が $4$ 日目の試験を終えて、あり得る限り最高の順位になるような $N$ 人の生徒の点の取り方を考えます。これは単純な話ですが、
- 生徒 $x$ は $300$ 点満点を取る
- それ以外の生徒は全員 $0$ 点を取る
のが、生徒 $x$ にとって最も都合の良い条件です。
各生徒 $i$ の $3$ 日目までの点数の合計を $S_i$ （$=P_{i,1} + P_{i,2} + P_{i,3}$）とします。$4$ 日目の試験の結果、生徒 $x$ は $S_x+300$ 点、他の各生徒 $i$ は $S_i$ 点を取ります。
このとき、生徒 $x$ が $K$ 位以内に入れる条件は、『$S_1$ ～ $S_N$ の中で $K$ 番目に大きい点数』を $T$ とおいて、$S_x + 300 \ge{T}$ であることです。$T$ 点ちょうどならば、最悪でも $K$ 位 ちょうど、最高の場合は $1$ 位になれるからです。（問題文に『$4$ 日目の試験後の生徒の順位は、その生徒よりも $4$ 日間の合計点が高い生徒の人数に $1$ を加えた値として定めます。』とあります）
実装
まず、$N$ 人それぞれの $3$ 日間の合計点数が入った配列 S を作ります。何日目に何点取ったという情報は不要なので、入力から合計を求めたら捨ててしまって構いません。
次に、S を降順（大きい順）にソートした配列 S_sorted を sorted(S)で作ります。（ソートされる前の 生徒番号順に並んだ S も必要です)
S_sorted から $K$ 番目に大きい値を取り出して、 $T$ とします。Pythonのインデックスは $0$ スタートですから、T = S_sorted[K-1] が $K$ 番目に大きい値です。もしくは、S_sorted の先頭にダミーの要素を $1$ つ追加して、S_sorted[x] が $x$ 番目に大きい値を表すようにズラして、T = S_sorted[K]としても良いです。
コード
N, K = map(int, input().split())
S = [sum(map(int, input().split())) for _ in range(N)] # 3日間の合計点数だけわかればいいです
S_sorted = [10000] + sorted(S, reverse=True) # S_sorted[K]が暫定K位を表すよう、0番目にダミーを入れます
border = S_sorted[K] # 暫定K位の点数です
for score in S:
print("Yes" if score + 300 >= border else "No") # 同点でもOKです
D問題『Linear Probing』
問題ページ：D - Linear Probing
灰コーダー正解率：3.9 %
茶コーダー正解率：20.3 %
緑コーダー正解率：49.7 %
入力
$Q$ : クエリの数
$t_i$ : クエリ $i$ の種類（$1$ か $2$）
$x_i$ : クエリで用いる整数（$0\le{x_i}\le{10^{18}}$）
入力では与えられない定数
$N=2^{20}=1048576$
注意
この問題の実行時間制限は $4$ 秒です。
考察
数列に対してクエリが与えられるので、それを処理してくださいという問題です。また、この問題では、$N=2^{20}=1048576$ （$100$ 万ちょっとです）という問題文で書いてある定数を使用します。
長さ $N$ の数列 $A=(A_0,A_1,\dots,A_{N-1})$ があって、はじめは全要素 $-1$ です。Pythonでいうと
N = 2 ** 20
A = [-1] * N
ということです。
クエリは、クエリ$1$ と クエリ $2$ の $2$ 種類あります。クエリ $2$ のほうが簡単なので、先にクエリ $2$ を説明します。
クエリ2
『その時点での $A_{x_i\ mod\ N}$ の値を出力する。』
$A$ の「$x_i$ を数列の長さ $N$ で割った余り」番目の値を出力するという意味です。（必ず $0$ 以上 $N$ 未満になります）
Pythonでいうと、以下の単純なことをするだけです。
print(A[x_i % N])
クエリ1
クエリ $1$ の操作内容をわかりやすく書くとこうなります。
- 整数 $h$ を $h=x_i\ mod\ N$ とする。
- $A_h = -1$ の要素が見つかるまで、$h$ を $1$ ずつ増やす。ただし、$h=N-1$ の次は $h=0$ とする。（この操作は無限ループにならない）
- $A_h$ の値を、$-1$ から $x_i$ に変更する。
$0\le{x_i}\le{10^{18}}$ ですから、一度書き換えられた要素はもう二度と $-1$ には戻りません。つまり、『まだ書き換えられていない番号で、$h$ の初期値から"最も近い"もの』を探せばいいです。（最も近いとは、$h$ を $1$ ずつ増やしていったとき、最初に見つかる書き換えられていない番号ということです）
問題文通りの操作をシミュレートすると、計算量が $O(QN)$ となりTLEになります。そのため、高速に『まだ書き換えられていない番号で"最も近い"もの』を見つける必要があります。
解法
この問題には様々な解法があります。
- 『経路圧縮』を用いる方法（データ構造のUnionFindでも使われているものです）
- 『セグメントツリーによるRange Minimum Query』を用いる方法
- 『順序付き集合』を用いる方法
- 『平方分割』を用いる方法
自分の実装の場合、解法 $1$ と $2$ はPython 3でACを確認しました。解法 $3$ と $4$ はPythonではTLEになるので、PyPyで出してください。良い実装であれば、解法 $3$ も余裕を持ってACできると思います。
『経路圧縮』を用いる方法
データ構造のUnionFindと似た構造で番号を管理し、高速に求めます。
長さ $N$ の 配列 $P = (P_0, P_1,\dots,P_{N-1})$ を用意し、初期値は $P_i=i$ とします。この配列は、『自分に最も近い、書き換えられていない番号を見つけるための補助配列』 です。$P_i = i$ は、その番号がまだ書き換えられていないことを示します。うまく $P$ を書き換えることで、$P$ の要素をたどって書き換えられていない番号を見つけようというアイデアです。
$A_h$ が $-1$ である要素が見つかるまで $h$ を $1$ ずつ増やす代わりに、$h\larr P_h$ と代入します。こうすることで、$1$ ずつ探索するよりも高速に $P_h=h$ である要素を探します。
ある $A_h$ が $x_i$ に書き換えられたとき、$P_h$ を $P_{h+1}$ に書きかえます。$P_h$ 自身はもう書き換えられてしまったので、$1$ つ隣の $P_{h+1}$ が指している番号を $P_h$ も指すようにします。しかし、これだけでは、ただシミュレーションを行っているのと全く同じことしているだけです。
そこで『経路圧縮』を行います。
- $A_h$ を書き換えたら $P_h$ だけでなく、スタートから辿っていった $P_i$ すべてを $P_{h+1}$ に書き換える
『経路圧縮』を用いることで、各操作をならし計算量 $O(logN)$ で行うことができます。
つまりどういうことですか
UnionFindです。
配列の番号を頂点とした、有向根付き木の森（複数の有向根付き木）を考えます。
$P$ は 有向グラフの辺を表します。初期状態ではどの頂点も書き換えられていませんから、$1$ 頂点の木が $N$ 個 存在しています。
木の根の頂点番号が、まだ書き換えられていない番号を表します。根でない頂点は、すでに書き換えられた頂点です。根でない頂点から辺を辿っていけば、最終的には根にたどり着いて、書き換えられていない番号がわかります。
しかし、何も工夫をせずに辺を貼っているだけでは、木の高さに偏りが生じる（とくにこの問題の場合直線上になる）ので、根までたどる頂点の数が多くなり、計算量が悪化します。
そこで、根を見つけるときに辿った頂点すべてについて、辺が直接新しい根を指すように貼り替えてしまいます。こうすることで、木の高さが高くならないように抑えられます。
『有向根付き木の森』構造や経路圧縮自体、UnionFindでやっていることと全く同じです。ただし、一般的なUnionFindはUnion by Rankという方法も併用して木の高さを抑えており、この問題で使うと根の番号が最も近い書き換えられていない番号を指さなくなる恐れがあります。
『セグメントツリーによるRange Minimum Query』を用いる方法
仕組みさえ理解できれば、この方法が一番楽だと思います。
Range Minimum Queryとは
Range Minimum Query（RMQ）とは、以下のクエリが可能なデータ構造のことです。
- 配列の区間を指定して、区間内で最小の値を求める
- 配列の $1$ 点の要素を変更する（これができないSparse TableなどもRMQと呼ばれますが）
計算量を気にしなければ、Python組み込みのリストとmin関数だけでも実現可能です。
A = [1, 5, 3, 4, 2]
m1 = min(A[0:3]) # == 1
m2 = min(A[1:5]) # == 2
A[3] = 0
m3 = min(A[2:4]) # == 0
もちろん、この方法は遅すぎて使い物になりません。
RMQの各操作を $O(logN)$ で行えるデータ構造に、セグメントツリーというものがあります。
RMQで解く方法
まだ書き換えられていない番号を、セグメントツリーで管理します。
下準備は以下の $2$ つです。
- RMQができる長さ $N$ の セグメントツリー
segを用意する - すべての
iについて、seg[i] = iに初期化する
クエリ $1$ で『書き換えるべき番号』を求めるときは、以下の操作を行います。なお、INF には $N$ 以上の適当な整数を使います。
- 区間 $[h, N)$ の最小値 $pos$ を求めて、
INFでなければを $A_{pos}=x_i$ に書き換える - 手順 $1$ で
INFが返ってきた場合、区間 $[0, h)$ の最小値 $pos$を求めて、$A_{pos}=x_i$ に書き換える seg[pos] = INFに変更する
インデックスそのものをセグ木にのせることで、$h$ から最も近い（=最小）のインデックスをRMQで求めることができます。すでに書き換えた番号をINFにすることで、RMQの最小値に引っかからないようにしています。
順序付き集合を用いる方法
順序付き集合というデータ構造で、まだ使っていない番号を管理します。
順序付き集合は「集合をソートされた順に保ったまま要素を $O(log(N))$ で追加・削除」できるデータ構造です。順序付き集合を二分探索することで、$h$ 以上でまだ使っていない番号を $O(log(N))$ で求めることができます。
C++では、標準ライブラリにstd::setという順序付き集合がありますが、Pythonにはありません。そのため、自分で実装する必要があります。一般的には、平衡二分探索木というデータ構造で実装されます。
自分の実装の計算量は $O(N\ log\ N + Q\ log\ N)$ ですが、公式解説の解法では $O(N+Q\ log\ Q)$ または $O(Q\ log\ Q)$ です。
平方分割を用いる方法
$N=2^{20}$ の配列を、区間ごとに長さ $\sqrt{N}$ の配列 $\sqrt{N}$ 個に分割します。$\sqrt{2^{20}}=2^{10}=1024$ ですから、長さ $1024$ の配列を $1024$ 個、二次元リストで管理します。
また、区間ごとに $-1$ の要素が残りいくつあるか管理する配列を作ります。初期値は $1024$ です。
クエリ $1$ では、以下の操作を行います。
- $h$ が属する区間の、$h$ 以降の要素に $-1$ があるかループで探す（$O(\sqrt{N})$）
- 1で見つからなければ、最も近い $-1$ の要素が $1$ 個以上ある区間をループで探す
- *見つけた番号の要素を $x_i$ で書き換え、その番号が属する区間の $-1$ の残り数カウントを1減らす *
計算量は $O(Q\sqrt{N})$ です。
コード
経路圧縮を使う方法
Pythonで通ります。（750ms）
def main():
N = 2 ** 20
Q = int(input())
A = [-1] * N
P = list(range(N))
for _ in range(Q):
t, x = map(int, input().split())
if t == 1:
h = x % N
pos = h
visited = [pos]
while A[pos] != -1:
pos = P[pos]
visited.append(pos)
A[pos] = x
new_p = P[(pos + 1) % N]
for u in visited:
P[u] = new_p
else:
print(A[x % N])
if __name__ == '__main__':
main()
セグメント木でRange Minimum Query（RMQ, 区間の最小を取得する）を使う方法
Pythonでもギリギリ通ります。（3500ms）良い実装ならもっと早くなると思います。
RMQセグ木のコード
class SegmentTreeRMQ:
def __init__(self, n, e=1 << 60, array=None):
self.n = n
self.e = e
self.log = (self.n - 1).bit_length()
self.size = 1 << self.log
self.d = [e] * (2 * self.size)
class SegmentTreeRMQ:
def __init__(self, n, e=1 << 60, array=None):
self.n = n
self.e = e
self.log = (self.n - 1).bit_length()
self.size = 1 << self.log
self.d = [e] * (2 * self.size)
if array:
for i in range(self.n):
self.d[self.size + i] = array[i]
for i in reversed(range(1, self.size)):
self.__update(i)
def set(self, p, x):
"""
a[p] に x を代入する
"""
p += self.size
self.d[p] = x
for i in range(1, self.log + 1):
self.__update(p >> i)
def get(self, p):
"""
a[p]を返す
"""
return self.d[p + self.size]
def prod(self, l, r):
"""
[l, r) の総積を返す
"""
sml = self.e
smr = self.e
l += self.size
r += self.size
while l < r:
if l & 1:
sml = min(sml, self.d[l])
l += 1
if r & 1:
r -= 1
smr = min(self.d[r], smr)
l >>= 1
r >>= 1
return min(sml, smr)
def all_prod(self):
"""
[0, n) の総積を返す
"""
return self.d[1]
def max_right(self, l, f):
if l == self.n:
return self.n
size = self.size
l += size
sm = self.e
while True:
while not (l & 1):
l >>= 1
if not f(min(sm, self.d[l])):
while l < size:
l <<= 1
if f(min(sm, self.d[l])):
sm = min(sm, self.d[l])
l += 1
return l - size
sm = min(sm, self.d[l])
l += 1
if (l & -l) == l:
break
return self.n
def min_left(self, r, f):
if r == 0:
return 0
size = self.size
r += self.size
sm = self.e
while True:
r -= 1
while r and r & 1:
r >>= 1
if not f(min(self.d[r], sm)):
while r < size:
r = 2 * r + 1
if f(min(self.d[r], sm)):
sm = min(self.d[r], sm)
r -= 1
return r + 1 - size
sm = min(self.d[r], sm)
if (r & -r) == r:
break
return 0
def __update(self, k):
self.d[k] = min(self.d[k << 1], self.d[k << 1 | 1])
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.n
return self.prod(start, stop)
else:
return self.d[key + self.size]
def __setitem__(self, key, value):
key += self.size
self.d[key] = value
for i in range(1, self.log + 1):
self.__update(key >> i)
def main():
def query1(h):
ind = seg[h:] # インデックスやスライスでセグメント木を操作できる実装にしています
if ind == INF:
ind = seg[:h]
seg[ind] = INF
return ind
INF = 1 << 30
N = 2 ** 20
Q = int(input())
A = [-1] * N
seg = SegmentTreeRMQ(N, INF, array=list(range(N)))
for _ in range(Q):
t, x = map(int, input().split())
h = x % N
if t == 1:
ind = query1(h)
A[ind] = x
else:
print(A[h])
if __name__ == '__main__':
main()
順序付き集合を使う方法
自分の実装は遅いので、Pythonでは通りませんでした。まともな順序付き集合の実装や、Binary Index Treeを使った順序付き集合であれば通ると思います。
平衡二分探索木のコード（遅いです）
class OrderedSet:
"""
平衡二分探索木（Treap）
"""
from typing import List, Optional, TypeVar
class TreapNode:
def __init__(self, key, priority):
self.key = key
self.left = None
self.right = None
self.priority = priority
def __repr__(self):
return str(self.key)
T = TypeVar('T', int, float, str)
Node = TreapNode
def __init__(self):
from random import randint
self.rand_gen = randint
self.root = self._create_node(None)
self.length = 0
self._none_node = self._create_node(None)
def is_empty(self):
if self.root is None:
return True
else:
return False
def _create_node(self, x) -> Node:
priority = self.rand_gen(1, (1 << 32) - 1)
new_node = self.TreapNode(x, priority)
return new_node
def add(self, x) -> None:
"""
xを追加する
既に含まれている場合は何もしない
"""
if self._add(x):
self.length += 1
def _add(self, x) -> bool:
"""
追加に成功した場合True、既に存在した場合Falseを返す(len管理用)
"""
if self.root.key is None:
self.root = self._create_node(x)
return True
cur = self.root
par_stack = []
while cur:
par_stack.append(cur)
if x < cur.key:
cur = cur.left
elif x > cur.key:
cur = cur.right
else:
return False
l = len(par_stack)
child = self._create_node(x)
for i in range(l):
par = par_stack[-(i + 1)]
if child.key < par.key:
par.left = child
if child.priority < par.priority:
child.right, par.left = par, child.right
par = child
else:
par.right = child
if child.priority < par.priority:
child.left, par.right = par, child.left
par = child
child = par
self.root = child
return True
def discard(self, x) -> None:
"""
xを削除
なければ何もしない
"""
if self._discard(x):
self.length -= 1
def _discard(self, x) -> bool:
if self.root.key is None:
return False
cur = self.root
node_stack = []
is_left_stack = []
while True:
if cur is None:
return False
if x < cur.key:
node_stack.append(cur)
is_left_stack.append(True)
cur = cur.left
elif x > cur.key:
node_stack.append(cur)
is_left_stack.append(False)
cur = cur.right
else:
if cur.left and cur.right:
if cur.left.priority < cur.right.priority:
node = cur.left
node.right, cur.left = cur, node.right
node_stack.append(node)
is_left_stack.append(False)
else:
node = cur.right
node.left, cur.right = cur, node.left
node_stack.append(node)
is_left_stack.append(True)
else:
node_stack.append(cur.left or cur.right) # Noneでないほうが追加される
is_left_stack.append(None)
break
l = len(node_stack)
cur = node_stack[-1]
is_left = is_left_stack[-1]
for i in range(1, l):
par = node_stack[-(i + 1)]
is_left = is_left_stack[-(i + 1)]
if is_left:
par.left = cur
else:
par.right = cur
cur = par
self.root = cur
return True
def find(self, x) -> bool:
"""
xが含まれるか判定
計算量: O(logN)
"""
if self._find(x):
return True
else:
return False
def _find(self, x) -> Optional[Node]:
if self.root.key is None:
return None
node = self.root
while node:
if x < node.key:
node = node.left
elif x == node.key:
return node
else:
node = node.right
return None
def all_keys(self) -> List[T]:
"""
全要素の昇順リスト
Todo: 軽量化
"""
def add_sub_tree(node):
keys = []
if node is not None:
keys.extend(add_sub_tree(node.left))
keys.append(node.key)
keys.extend(add_sub_tree(node.right))
return keys
if self.length == 0:
return []
else:
keys = add_sub_tree(self.root)
return keys
def min(self) -> T:
node = self.root
while node.left:
node = node.left
return node.key
def max(self) -> T:
node = self.root
while node.right:
node = node.right
return node.key
def prev(self, x) -> Optional[T]:
"""
xのより真に小さい要素で最大のものを取得
存在しなければNoneを返す
計算量: O(logN)
"""
if self.root.key is None:
return None
node_cur = self.root
node_return = self._none_node
while True:
if node_cur.key < x:
node_return = node_cur
if node_cur.right:
node_cur = node_cur.right
else:
return node_return.key
else:
if node_cur.left:
node_cur = node_cur.left
else:
return node_return.key
def next(self, x) -> Optional[T]:
"""
xより真に大きい要素で最小のものを取得
存在しなければNoneを返す
計算量: O(logN)
"""
if self.root.key is None:
return None
node_cur = self.root
node_return = self._none_node
while True:
if x < node_cur.key:
node_return = node_cur
if node_cur.left:
node_cur = node_cur.left
else:
return node_return.key
else:
if node_cur.right:
node_cur = node_cur.right
else:
return node_return.key
def __str__(self):
return str(self.all_keys())
def __repr__(self):
return self.__str__()
def __len__(self):
return self.length
def __contains__(self, x):
return self.find(x)
def main():
def query1():
h = x % N
if A[h] == -1:
return h
nxt = ods.next(h) # hより大きい要素で最小のものを返すメソッドです（"hの次"を得られます）
if nxt is not None:
return nxt
else:
# hより大きい要素がない（＝hが最大の要素）場合は、0以上で最小の要素を探します
if A[0] == -1:
return 0
return ods.next(0)
N = 2 ** 20 # 問題文にある定数です
Q = int(input())
A = [-1] * N # 数列Aそのものです
ods = OrderedSet() # A[i] == -1であるiを順序付き集合で管理します
for i in range(N):
ods.add(i) # はじめは全要素-1なので、0～Nまで順序付き集合に追加します
for _ in range(Q):
t, x = map(int, input().split())
if t == 1:
pos = query1(x)
A[pos] = x
ods.discard(pos) # posが-1ではなくなったので順序付き集合から削除します
else:
print(A[x % N])
if __name__ == '__main__':
main()
平方分割を使う方法
分割の区間数を変えるなど色々がんばりましたが、Pythonでは1ケースだけTLEが取れませんでした。
def main():
def query1(fi, fj):
# A[fi][fj:]に-1があるか探索
if cnt[fi] > 0:
for j in range(fj, SQ):
if A[fi][j] == -1:
A[fi][j] = x
cnt[fi] -= 1
return
# A[(fi + c)]に-1があるか探索
for c in range(1, SQ):
i = (fi + c) % SQ
if cnt[i] == 0: # -1がない区間は飛ばさないと、全探索と変わりません
continue
for j in range(SQ):
if A[i][j] == -1:
A[i][j] = x
cnt[i] -= 1
return
N = 2 ** 20
SQ = 2 ** 10 # √N = 1024
Q = int(input())
A = [[-1] * SQ for _ in range(SQ)] # A[i][j]: 1024i + j
cnt = [SQ] * SQ # cnt[i]: 1024i～1024(i+1)-1の区間に、-1の要素がいくつあるか（0ならスキップする）
for _ in range(Q):
t, x = map(int, input().split())
h = x % N
fi, fj = h // SQ, h % SQ
if t == 1:
query1(fi, fj)
else:
print(A[fi][fj])
if __name__ == '__main__':
main()