AtCoder Beginner Contest 097に事後参加してみた
とりあえずこれから何週間か、1日1つ以上ABCの4問題を解いていこうと思っている。
D問題は読んでから少し放置して暇なときにつらつら考えて、半日たってから実装したりしているので全然コンテスト形式は再現できていないのだが、とりあえず今は自分で問題に対していろいろとアプローチしていくことで経験と勘を育てていきたい。
というわけで今回はABC097
第１問
直線上のABCの点のうち、AとCがBを経由して・あるいは経由せずに通信できるか。ただし二点は距離がD以内でないと通信できない。
a,b,c,d = map(int, input().split()) print('Yes' if abs(a-c) <= d or all(abs(x-b) <= d for x in (a,c)) else 'No')
builtinのallが地味に便利。
第２問
X以下の最大のベキ乗数を求める。
from itertools import chain, count, takewhile x = int(input()) if x > 1: ps = takewhile(lambda p: 2**p <= x, count(2)) bps = (takewhile(lambda y:y <= x, (b**p for b in count(1))) for p in ps) print(max(chain.from_iterable(bps))) else: print(1)
xが1以下だと冪指数のシーケンスが空になってしまうので特殊ケース化している。できれば避けたいのだがうまく解決する方法が思い浮かばなかった。
第３問
今回の中でちょっと満足しているコード。
import heapq from itertools import groupby, islice s = input() n = int(input()) def substrings_starting_at(i): return (s[i:j+1] for j in range(i, len(s))) gens = [substrings_starting_at(i) for i in range(len(s))] sm = heapq.merge(*gens) res, _ = next(islice(groupby(sm), n-1, n)) print(res)
文字列sの先頭から始まる部分文字列はs[:1] < s[:2] < s[:3] .. <sであることを利用している。
つまり[s[n:n+1], s[n:n+2], s[n:n+3]]はすでにソート済みだということ。
なので[s[0:1], s[0:2] ... s[0:]], [s[1:2], s[1:3], ... s[1:]] ...]をmerge sortと同じ要領でmergeしてしまえばいい。heapq.mergeでそれができる。
uniqueな文字列のみという制限があるので、ソートされていることを利用してgroupbyで同一の要素は集めてしまって、あとはn番目の要素を取り出しているだけ。
解法とPythonの機能利用の両側面からそれなりに上手くできたように思っている。
第４問
1..Nの数字が並び変えられた配列のうち、指定された複数の添字ペアだけ何回でも内容をスワップすることができる。配列の添字と数字が一致するセルの個数を最大化する。
aとb、bとcがお互いにスワップ可能な場合、abcをどのような配置にも入れ替え可能になる。なので「相互につながっていてスワップ可能なセル群」を探して、「ある数字が入っているセルがその数字の添字のセルとつながっている」回数を数えればいい。
n, m = map(int, input().split()) ps = [int(p)-1 for p in input().split()] xys = [tuple(int(x)-1 for x in input().split()) for _ in range(m)] union = {node:node for node in range(n)} def root(node): if union[node] == node: return node union[node] = root(union[node]) return union[node] for x, y in xys: union[root(x)] = root(y) print(sum(root(i) == root(p) for i, p in enumerate(ps)))
最低限のUnion Find木。経路圧縮はしているがランクによる結合順の選定はしていない。なのでasymptoticな計算量はO(nlogn)なのだが、実行してみるとテストケースではランクを考慮に入れた実装より早かった。
上記の通り、つながっているかどうかを調べて数える、ということをそのまんまでやっているだけ。経路圧縮Union Find木は辞書と再帰で実装が楽だ。
感想
heapq.mergeやgroupby、union find木など知ってはいたけれどもあまり引っ張り出して使ったことがない機能をいろいろ組み合わせたりできるのはかなり楽しい。これらを自然と使えるようになればコンテストの時間制限の中でもけっこう効率的できれいな解法に行き着きやすいのではないかと淡い期待を抱いている。
