蟻本初級編攻略 - 2-1 特殊な状態の列挙
蟻本ではちゃんと例題めいたものが載っていなかった話題。
Pythonだとitertools.permutationなどで簡単に実装できる。
ABC054C - One-stroke Path
重み無し無向グラフを、特定の始点「1」からはじめてすべての頂点を一回だけ訪れるパスを数える問題。
「1」以外の頂点を並べる順番をすべて列挙してから始点に「1」を加え、隣同士のペアがすべて辺の集合に含まれているかを調べる。
from itertools import permutations n, m = map(int, input().split()) edges = set() for _ in range(m): a, b = map(int, input().split()) edges.add((a, b)) edges.add((b, a)) orderings = ((1,) + p for p in permutations(range(2, n+1))) print(sum(all(move in edges for move in zip(p, p[1:])) for p in orderings))
ちなみに始点「1」を加えるのに、AtCoderで使えるPython3のバージョンが3.4なので(1,)+p
としている。
Python3.5からはUnpacking SyntaxがPEP448で拡張されて'(1, *p)'と書ける。
蟻本初級編攻略 - 2-1 迷路の最短路
ここらへんは全部今年の三月に記事を書いているなー。
このころは
map_をdict化すればかなり綺麗になるが、競プロ的にどうなんだろう・・・
などと言っていたが、現在は躊躇なくdictionary化するなあ。実際そこのO(N)が問題になることはまずない。(すくなくともABCのD問題くらいでは)
「もしコストが気になるならはじめからdictionaryとしてパースしてしまえばいいじゃない」と心の中のIQ145の女子高生に囁かれて試してみたのが以下のコード:
from collections import deque n, m = map(int, input().split()) maze = {(i,j):c for i in range(n) for j, c in enumerate(input())} def bfs(coord): dq = deque() dq.appendleft((0, coord)) while dq: steps, (i, j) = dq.pop() current = maze.get((i, j), '#') if current == 'G': return steps if current == '#': continue maze[(i, j)] = '#' dirs = ((-1, 0), (1, 0), (0, 1), (0, -1)) dq.extendleft((steps+1, (i+di, j+dj)) for di, dj in dirs) start = next(coord for coord, c in maze.items() if c == 'S') print(bfs(start))
DFSをスタックでやることの隠れたメリットはBFSとDFSがほぼ同型になること。DFSではLIFOのスタックを使っていたのをBFSではFIFOのキューを使う、以外はほぼ同じコード。
ABC007C - 幅優先探索
蟻本の問題とほぼ同じ。スタートとゴールが地図とは別に座標で与えられることとその座標が0-indexedじゃないことが注意点。
from collections import deque n, m = map(int, input().split()) start = tuple(int(x)-1 for x in input().split()) goal = tuple(int(x)-1 for x in input().split()) maze = {(i,j):c for i in range(n) for j, c in enumerate(input())} def bfs(coord): dq = deque() dq.appendleft((0, coord)) while dq: steps, (i, j) = dq.pop() if (i, j) == goal: return steps if maze.get((i, j), '#') == '#': continue maze[(i, j)] = '#' dirs = ((-1, 0), (1, 0), (0, 1), (0, -1)) dq.extendleft((steps+1, (i+di, j+dj)) for di, dj in dirs) print(bfs(start))
AtCoderはほぼすべての問題が1-indexedで修正するのをうっかりすると後から見つけにくい。
ABC088D - Grid Repainting
「BFSの結果のルート」と「元からあった黒いマス」以外のマスをすべて黒く塗れるという点に気づけば、あとはほぼ同じ問題。
from collections import deque h, w = map(int, input().split()) maze = {(i,j):c for i in range(h) for j, c in enumerate(input())} blacks = sum(c == '#' for c in maze.values()) def bfs(i, j): dq = deque() dq.appendleft((1, (i, j))) while dq: steps, (i, j) = dq.pop() if (i, j) == (h-1, w-1): return steps if maze.get((i, j), '#') == '#': continue maze[(i, j)] = '#' dirs = ((-1, 0), (1, 0), (0, 1), (0, -1)) dq.extendleft((steps+1, (i+di, j+dj)) for di, dj in dirs) res = bfs(0, 0) print(-1 if res is None else h*w - res - blacks)
ただ、この問題ではそもそもスタート地点からゴールまでの経路が存在しないケースがあり得る。最初に提出したときはうっかりしていた・・・
ARC005C - 器物損壊!高橋君
壁を二回まで破壊して通過できるという条件でスタートからゴールまで到達できるか、という問題。
最初は「DFSで隣接している壁を探して、隣接している壁がダブっているかどうかで繋がっている『部屋』を見つけて2ステップでゴールまで到達できるか」的なコードを書いていたのだが、あえなくTLEした。
その後に記事の説明を読んで01-BFSなる用語を調べたところ、非常にエレガントな解法で感動した。
「距離」を「移動コスト」再定義して「通路の移動コストは0」「壁の移動コストは1」とすれば「壁を最大二回移動してゴールに到達できるか」が算出できる。
h, w = map(int, input().split()) maze = {(i,j):c for i in range(h) for j, c in enumerate(input())} dirs = ((-1, 0), (1, 0), (0, 1), (0, -1)) def bfs01(start, maze=maze, dirs=dirs): q = [[start], [], []] while q: i, j = q[0].pop() maze[i, j] = 'X' for di, dj in dirs: next_ = maze.get((i+di, j+dj), 'X') if next_ == 'g': return True if next_ == 'X' or (next_ == '#' and len(q) == 1): continue q[next_ == '#'].append((i+di, j+dj)) while q and not q[0]: q = q[1:] return False start = next(coord for coord, c in maze.items() if c == 's') print('YES' if bfs01(start) else 'NO')
ただ、アイディアのエレガンスの対してコード実装が泥臭いのが難点・・・
最初はPriority QueueでやってみたのだがやはりO(n log n)だとTLEになる。リストを束ねて「スタックのキュー」的なデータ構造にした。
これで実行時間は800ms弱。Pythonだとここらへんが上限だろうか?
ちなみにPyPy3で同じコードを走らせてみると、Python3で20msぐらいかかるテストケースで165msかかり、すこし重いテストケースは軒並みMLEを起こしていた。
すくなくともAtCoderでPyPy3は使い物にならないと言ってしまってもいいのではないか。
理由としては:
- JITコンパイラが温まりきらない(数秒は継続して走っていないとJITの恩恵を受けない)
- PyPyはCPythonよりスタートアップに時間がかかる
- PyPy3はPyPy2に比べて最適化などがまだ甘い(Python2.7の仕様がもう10年もほとんど変わっていないので追いやすかった?)
- AtCoderが使っているPyPy3のバージョンが「CPython3より遅いことも多い」と知られている古いバージョン
- PyPyはitertoolsやlist comprehensionなどよりfor loopなどのほうが効率化しやすい
などが挙げられると思う。これらを合わせると正直PyPy3は当面無視するのが得策だろう。AtCoderでCythonが使えたらなかなか面白いことになるとは思うが・・・
蟻本初級編攻略 - 2-1 Lake Counting
やはり以前もブログに記事を書いた問題。
やはりコードが長い、というかごちゃごちゃしている・・・
今だったらこう書く:
from itertools import product n, m = map(int, input().split()) lake = [input() for _ in range(n)] lakedict = {(i, j):c for i, row in enumerate(lake) for j, c in enumerate(row)} def dfs(i, j): if lakedict.get((i, j), '.') != 'W': return lakedict[i, j] = '.' for di, dj in product((-1, 0, 1), repeat=2): dfs(i+di, j+dj) n = 0 for i, j in lakedict: if lakedict[i, j] == 'W': n += 1 dfs(i, j) print(n)
ATC001A - 深さ優先探索
失敗1:Recursion Limit
まずは非常に素直にDFSを実装してみた:
from itertools import product h, w = map(int, input().split()) xs = [input() for _ in range(h)] mapdict = {(i, j):c for i, row in enumerate(xs) for j, c in enumerate(row)} def dfs(i, j): current = mapdict.get((i, j), '#') if current == '#': return False if current == 'g': return True mapdict[i, j] = '#' dirs = ((-1, 0), (1, 0), (0, 1), (0, -1)) return any(dfs(i+di, j+dj) for di, dj in dirs) start = next(coord for coord, c in mapdict.items() if c == 's') print('Yes' if dfs(*start) else 'No')
サンプルケースは全部通るし、コード自体も見通しは悪くないのでバグはなさそう。ということで走らせてみると見事に複数のRE。再帰リミットにぶつかっているようだ。
from sys import setrecursionlimit setrecursionlimit(5000000)
というわけで再帰リミットをあげてみた。これでほとんど通るのだが、ひとつだけMLE。
Pythonでは非常に深い再帰の問題は再帰リミットだけではなく、関数スタックが大きいのでメモリ使用が厳しいという点もある。末尾呼び出し最適化もないし。
というわけで、まあ再帰しなければいいよね。という単純な解決。
from itertools import product h, w = map(int, input().split()) xs = [input() for _ in range(h)] mapdict = {(i, j):c for i, row in enumerate(xs) for j, c in enumerate(row)} def dfs(start): stack = [start] while stack: i, j = stack.pop() current = mapdict.get((i, j), '#') if current == '#': continue if current == 'g': return True mapdict[i, j] = '#' dirs = ((-1, 0), (1, 0), (0, 1), (0, -1)) stack.extend((i+di, j+dj) for di, dj in dirs) return False start = next(coord for coord, c in mapdict.items() if c == 's') print('Yes' if dfs(start) else 'No')
この問題では可能か否かを葉ノードで判定してそのブール値をそのまま木の上まで返しているだけなので、スタックを使ってしまえばあまりコードをいじらずに再帰を避けることができる。
これは無事にAC。
ARC031B - 埋め立て
似たような問題と似たような解法。
from functools import reduce from operator import and_ mapdict = {(i, j): c for i in range(10) for j, c in enumerate(input())} def dfs(start): stack = [start] while stack: i, j = stack.pop() current = mapdict.get((i, j), 'v') if current == 'v': continue if current == 'x': yield (i, j) continue mapdict[i, j] = 'v' dirs = ((-1, 0), (1, 0), (0, 1), (0, -1)) stack.extend((i+di, j+dj) for di, dj in dirs) surroundings = [] for coord in mapdict: if mapdict[coord] == 'o': surroundings.append(set(dfs(coord))) print('YES' if len(reduce(and_, surroundings)) else 'NO')
yieldを使って島の周囲の海の座標をシーケンスとして返している。それをset化し、「各島の周りの海の座標のset」すべてに共通する要素があるかを判定する。
ARC037B - バウムテスト
無指向グラフにいくつ非巡回グラフが含まれているかを数える問題。
問題文をみると???だが実際に解法は非常に似通っている:
from collections import defaultdict n, m = map(int, input().split()) graph = defaultdict(list) for _ in range(m): u, v = map(int, input().split()) graph[u].append(v) graph[v].append(u) visited = set() def dfs(v): stack = [(None, v)] while stack: last, current = stack.pop() if current in visited: return 0 visited.add(current) stack.extend((current, node) for node in graph[current] if node != last) return 1 ans = 0 for node in range(1, n+1): if node not in visited: ans += dfs(node) print(ans)
無指向グラフでの巡回を探知する時に、「たった今きたノードに戻れる事実は『巡回』として認識しない」というロジックを組み込むのが一番めんどくさかった。スタックに今のノードと次に調べるべきノードの両方をtupleとして乗せることで対応している。
蟻本初級編攻略 - 2-1 部分和
定期的に「蟻本を全部読むぞ!」と思って読み始めては放置、を繰り返している。
そもそもAtCoderやCodeForceに参加することなく「ある程度読んでから競プロ試してみよう」と思っていたのが間違いだったのではないか。
ABC100にリアルタイム参加もしたし、今度こそ蟻本を攻略したい。
幸いこういう素晴らしい記事もある:
本を読み、AtCoderの問題を解き、その解法をブログに書いていこうと思う。
部分和問題
蟻本の冒頭あたりについては今年の三月にも記事を書いていた。
というわけで書いたブログとコードを引っ張り出してみる:
が、なんだこれは・・・ ひどいねコレは。Pythonで添字をあまり使いたくないのはわかるがそれでこんなにコードが長く読みにくくなっては本末転倒も甚だしい。
と、完全に「数ヶ月前の自分のコードはひどく見える」を実践してしまった。
蟻本の題意である「Yes/Noを返す」を素直に実装するなら:
n = int(input()) xs = [int(x) for x in input().split()] k = int(input()) def dfs(i, total): if i == n: return k == total return dfs(i+1, total+xs[i]) or dfs(i+1, total) print('Yes' if dfs(0, 0) else 'No')
if i == n
をif i == n or k <= total
として途中で抜け出すような最適化も可能なことが一見してわかる。
もし「条件を満たす集合を返す」という実装をするとしてもこんな感じだろう:
n = int(input()) xs = [int(x) for x in input().split()] k = int(input()) def dfs(i, total): if k == total: yield [] if i == n or k < total: return for res in dfs(i+1, total+xs[i]): res.append(i) yield res for res in dfs(i+1, total): yield res res = list(dfs(0, 0)) sum_string = ' + '.join(str(xs[i]) for i in reversed(res[0])) if res else '' print('Yes ({} = {})'.format(k, sum_string) if res else 'No')
ちなみにPythonで枝刈りなしの全列挙ならitertoolsを使ったほうがいい:
from itertools import compress, product n = int(input()) xs = [int(x) for x in input().split()] k = int(input()) subsets = (list(compress(xs, selectors)) for selectors in product((0,1), repeat=n)) res = [subset for subset in subsets if sum(subset) == k] sum_string = ' + '.join(str(xs[i]) for i in reversed(res[0])) if res else '' print('Yes ({} = {})'.format(k, sum_string) if res else 'No')
再帰リミットにぶつかる心配をする必要もないし、より抽象度の高い概念で記述できる。
ABC061C - たくさんの数式
というわけでAtCoderの類題を解いてみる。
これは枝刈りなしの全列挙なのでitertoolsを使う:
from itertools import product s = input() s2 = '{}'.join(s) betweenss = product(('', '+'), repeat=len(s)-1) print(sum(eval(s2.format(*betweens)) for betweens in betweenss))
sの各文字の間に{}
を埋めて、そこに''と'+'のセットをあてはめて評価していく。
ABC079C - Train Ticket
ロジックはまったく同じ:
from itertools import product s = input() s2 = '{}'.join(s) opss = product('+-', repeat=3) lhss = (s2.format(*ops) for ops in opss) eqs = (lhs+'=7' for lhs in lhss if eval(lhs)==7) print(next(eqs))
「かならず解がひとつ以上存在する」「複数解がある場合はどれを返してもいい」という条件なのでジェネレータを作ってnextで最初の要素を取り出している。
AtCoder Beginner Contest 095に事後参加してみた
この回はD問題がけっこう面白かった。ちょっとコードも解説も多めかもしれない・・・
第1問
700円のラーメンに味玉、チャーシュー、ネギのトッピング(各100円)するか否かを示す'o'か'x'かが三つ並んだ文字列を与えられて、必要な金額を計算する問題。
print(700 + 100*input().count('o'))
そのまんま。
第2問
ドーナッツの種類がN個あり、種類ごとに必要な粉の量mがあり、全種類のドーナッツを少なくとも一つは作る必要があるという制約のもと、粉の量Xで最大いくつのドーナッツを作ることができるか、という問題。
n, x = map(int, input().split()) ms = [int(input()) for _ in range(n)] print(n + (x - sum(ms))//min(ms))
全種類作るということでN個のドーナッツをmの和の量の粉で作って、残った粉でもっともコストの低いドーナッツを最大限作った時の数を返している。
第3問
二種類のピザA、Bと、その二種類が半分ずつ合体したピザABが注文できる。AやBがほしい場合、AやBをそのまま買うか、ABを二枚買って半分にわけて一枚のAと一枚のBにする方法がある。AをX枚、Bを Y枚ほしい場合に一番安く買える金額を求める。
決定するべき分岐点は二つある。
まず「AとBのどちらも必要な数(つまりXとYの最小値)を個別で買って揃えるか、ABを買って揃えるか」でそれは(a+b)と(c*2)のどちらが安いかで決まる。
次に「残ったAかBの必要な数を個別で買うか、ABを買って余った部分は捨てるか」でそれはaと(c2)、あるいはbと(c2)のどちらが安いかで決まる。
ちなみにZがXとYの最小値として、(X-Z)と(Y-Z)のいずれかは0。
というロジックをコードにするとこうなる:
a, b, c, x, y = map(int, input().split()) z = min(x, y) total = z * min(a + b, 2 * c) total += (x - z) * min(a, 2 * c) total += (y - z) * min(b, 2 * c) print(total)
第4問
円形のカウンターに置いてある寿司に、スタート地点からの時計回りでの距離(=到達までのカロリー消費量)と寿司を食べた時に得られるカロリーの二つの値がある。時計回り、反時計回りに動き回って得られるカロリー値を最大化する問題。
時計回り、反時計回りの計算をなるべく対称にするため、「スタート地点からの時計回りでの距離のリスト」を「点と点の間の距離のリスト」に変換しておく。
ds = [j - i for i, j in zip(chain([0], xs), chain(xs, [c]))]
xsが各点の起点からの距離、dsが起点を最初と最後に入れた上での、各点どうしの距離のリスト。cはカウンターの全長。
一方向しか動かない縛り
手始めに問題を単純化して、どちらかの方向にしか進めない場合を考えてみる。
時計回りに動いた場合、各点まで到達することで得られるカロリーはその点までの寿司のカロリーの合計からその点までの距離の合計を引いたものだ。
pythonでaccumulate関数が「リストの添字nまでの暫定合計」を返してくれるので、それを利用するとこう書ける:
clockwise = chain([0], (v - d for v, d in zip(accumulate(vs), accumulate(ds))))
(vsは各点に置かれている寿司のカロリー値)
とりあえず「全く動かない場合に得られるカロリー」ということでシーケンスの先頭に0をいれておいた。
これを時計回り・反時計回りでやって最大値を取れば、一方向に進むだけの場合の最大カロリーがわかる。
clockwise = chain([0], (v - d for v, d in zip(accumulate(vs), accumulate(ds)))) anticlockwise = chain([0], (v - d for v, d in zip(accumulate(vs[::-1]), accumulate(ds[::-1])))) maximum = max(chain(clockwise, anticlockwise))
chainでイテレータを二つつなげて、maxでその最大値をとっている。
方向転換も可能な場合
上記を踏まえて、方向転換について考えてみる。
少し考えると、二回以上方向転換をする意味がないことがわかる。
二回方向転換しても二度目の方向転換の後は「すでに寿司をとってしまった距離を再度歩いている」だけなので、「時計回り・反時計回り・時計回り」といった進みかたをするより「反時計回り・時計回り」と進んだほうが同じ寿司をより少ない距離でとれる。
なので方向転換を考慮に入れなければいけないのは一回のみ。
さらに「方向転換しない」というのは「まず一方向に0距離進んでから逆方向に転換する」という「一回方向転換」の特殊ケースとして扱える。
ということでこの問題は「まず時計回りにすすんでから方向転換して反時計回り」「まず反時計回りにすすんでから方向転換して時計回り」の二つのケースを調べるだけで済む。
m点まで時計回りで進んで、起点に戻って、反時計回りで起点からm+1点まででカロリーが最大化される点に進む
あるいは
m点まで反時計回りで進んで、起点に戻って、時計回りで起点からm-1点まででカロリーが最大化される点に進む
の二ケースだ。
時計回りから反時計回りに転換するケースを考えよう。
まず時計回りにmの地点まで進む場合、得られるカロリーはmまでの寿司のカロリーの合計。使われるカロリーはmまでの距離の二倍だ。なぜ二倍かというと、「起点まで戻る」というコストもかかるから。
clockwise1 = chain([0], (v - d*2 for v, d in zip(accumulate(vs), accumulate(ds))))
いったん起点に戻ってから反時計回りに進む場合の得られるカロリーは前出の一方向のケース:
anticlockwise1 = chain([0], (v - d for v, d in zip(accumulate(vs[::-1]), accumulate(ds[::-1]))))
では時計回りにmの地点まで進んでから反時計回りに進んで得られる最大のカロリーは?となると「反時計回りで起点からm+1点まででカロリーが最大化される点に進む」必要が出てくる。anticlockwise1の最大値を求めるだけではだめだ。起点nからm+1までの[n, n-1, n-2, n-3 ..., m+1]における暫定最大値を計算する必要がある。
というわけでまた出てくるのがaccumulate。デフォルトでは暫定合計をリスト化するが、関数fを渡すことで暫定reduce(xs, f)をリスト化してくれる。なのでmax関数を渡すことで、[max(xs[:1]), max(xs[:2]), max(xs[:3]) .., max(xs[:n])]というリストをO(len(xs))で作成する。
accumulate(xs, func=max)を使って「m地点まで時計回りで進んでから起点に戻って反時計回りに進んだ場合のカロリー最大値」のリストを作成するpythonコード:
clock_anticlock = (x + y for x, y in zip(clockwise1, reversed(list(accumulate(anticlockwise1, func=max)))))
これを反時計回りでもやれば問題解決。
ということで全部のコード:(コード部分横スクロールできます)
from itertools import accumulate, chain n, c = map(int, input().split()) xs, vs = zip(*(map(int, input().split()) for _ in range(n))) ds = [j - i for i, j in zip(chain([0], xs), chain(xs, [c]))] clockwise1 = chain([0], (v - d*2 for v, d in zip(accumulate(vs), accumulate(ds)))) anticlockwise1 = chain([0], (v - d for v, d in zip(accumulate(vs[::-1]), accumulate(ds[::-1])))) clock_anticlock = (x + y for x, y in zip(clockwise1, reversed(list(accumulate(anticlockwise1, func=max))))) anticlockwise2 = chain([0], (v - d*2 for v, d in zip(accumulate(vs[::-1]), accumulate(ds[::-1])))) clockwise2 = chain([0], (v - d for v, d in zip(accumulate(vs), accumulate(ds)))) anticlock_clock = (x + y for x, y in zip(anticlockwise2, reversed(list(accumulate(clockwise2, func=max))))) print(max(chain(clock_anticlock, anticlock_clock)))
とりあえずO(n)。比較的複雑そうなロジックがzip, accumulate, chainを使って効率良く記述できるのは嬉しい。
感想
コードを書くよりブログを書くほうが何倍も難しかった・・・ 言語化って大変だな。
AtCoder Beginner Contest 100に参加してみた
初リアルタイム参加。言語はPython3。
「ABCって二週間おきにあるのかな?次は6/23か」とのほほんとしていたら
今日は記念すべき100回目のBeginner Contest、AtCoder Beginner Contest 100なので、ぜひぜひ参加してみてくださいな!
— chokudai(高橋 直大) (@chokudai) 2018年6月16日
Hello Worldが出力出来れば参加してみておっけーなレベルだよ!!!https://t.co/bCiCjMD2hT
こんなツイートが出ていて地味に驚く。
あと一週間あることを前提に過去問を最近のやつから順に解いていたのだが、最近出題されたものばっかり見てたということでなんとなく不利な気がしてしまう。(最近出題した問題と類似したものはすぐには出しにくいのでは?という考えから)
まあ仕方ないから当たって砕けろ、で初参加。
とりあえず30分弱でミスなし全完で順位は86位。実力からすると完全にできすぎでビギナーズラック以外の何物でもない・・・
めでたく初回ランク上昇上限の400点をもらい茶色コーダーに。
第1問
ケーキを16等分に切って、同じ人が隣り合わせのピースを取れないとして、二人が申告した数のピースを食べられるかどうかを判定する。
16等分と偶数で、ケーキをとる人数も2なので非常に楽。隣り合わせのピースはとれないのだから一人当たり最大で8しかとれず、お互いにとった数は干渉し合わない。
a, b = map(int, input().split()) print('Yay!' if a < 9 and b < 9 else ':(')
第2問
100でd={0, 1, 2}回割り切れる、小さい順でn個目の正の数字を出力する。
順位表見ているとけっこうつまづきポイントだったようだ。
from itertools import count, islice d, n = map(int, input().split()) gen = (i * (100**d) for i in count(1) if i % 100 != 0) print(next(islice(gen, n-1, n)))
このif i % 100 != 0
の部分を見落とす人が多かったのだろうか。
next(islice(gen, n-1, n))
はnth(gen, n)
的なイディオム。ちょっとごちゃついていて好みではないが・・・
第3問
黒板に書いてある数字を全部「3するか/2するかのどちらかで、すべてを3することはできず、/2の結果は整数でなければいけない」という処理を何回行えるか。
n = int(input()) xs = [int(x) for x in input().split()] def div2(x): while x % 2 == 0: x /= 2 yield 1 print(sum(sum(div2(x)) for x in xs))
停止条件は2で割る数字がなくなった時。2と3が互いに素なので、*3の処理は無視できる。
なので各数字が何回2で割り切れるか、の和。
div2
関数はもうちょっときれいに書けそう。
第4問
N個のケーキからM個選ぶ問題。各ケーキには三つの指標があって、各指標ごとの合計の絶対値の和を最大化させたい。
各指標の絶対値の合計でも、和の絶対値でもないのがポイント。
from itertools import product n, m = map(int, input().split()) xyzs = [[int(x) for x in input().split()] for _ in range(n)] multipliers = product((-1, 1), repeat=3) def max_sum(multiplier): sums = (sum(a*mul for a, mul in zip(xyz, multiplier)) for xyz in xyzs) top_m = sorted(sums, reverse=True)[:m] return sum(top_m) print(max(max_sum(multiplier) for multiplier in multipliers))
ただの合計の最大値を求める問題ならO(n log n)でn<=1000なので余裕で間に合う。
絶対値をとろうとする代わりに、各指標に+1か-1を掛け合わせて足したものの合計の最大値を8通り計算してその最大値をとればいい。
感想
AとB問題をやってるときが一番緊張した。初回は結果が良くてもレートに反映されにくいし、AとBさえノーミスだったらCとDはゆっくりでいいや、と考えていたせいだろう。なのでけっこういい精神状態でCとDに取り組めてラッキーだったように思う。
処理時間はCが最大で100ms弱、あとは20~30msくらい。やはりPythonでもABCだったらアルゴリズムが正しければ余裕で通る。そして標準ライブラリやiteratorプロトコルをうまく組み合わせるとけっこうきれいに、宣言的に記述できるように思う。
すごく楽しかった。次回もまた参加したい。
AtCoder Beginner Contest 096に事後参加してみた
今日も今日とてAtCoder ABCを解く。
第1問
1月1日、2月2日、10月10日といったようにX月X日のようになっている日が年始からa月b日までに何日あったかを数える問題。
a, b = map(int, input().split()) print(a-(a>b))
ブール値が1と0に評価されることを利用している。各月に一回月と日が一致する日があり、その日を過ぎているかどうかをa > bで判定している。
第2問
a, b, cの整数のいずれかを二倍に変えるという操作をk回行い、和を最大化する問題。
a, b, c = sorted(map(int, input().split())) k = int(input()) print(a+b+(c*(2**k)))
当然一番でかい数字をk回二倍にして足し合わせるのが最大。
第3問
隣接した二つのマスを黒く塗りつぶすという操作のみで、特定の白黒のパターンを再現できるか判定する問題。
h, w = map(int, input().split()) s = [input() for _ in range(h)] d = {(i, j):c for i, row in enumerate(s) for j, c in enumerate(row)} def next_to_black(i, j): directions = ((-1, 0), (1, 0), (0, -1), (0, 1)) surroundings = ((i + a, j + b) for a, b in directions) colors = (d.get(cell, '.') for cell in surroundings) return any(color == '#' for color in colors) black_coords = (coord for coord, c in d.items() if c == '#') paintable = all(next_to_black(*coord) for coord in black_coords) print('Yes' if paintable else 'No')
塗るべきパターンの黒いマスを列挙して、その周りの四つのマスのいずれかが黒いか判定する、というロジックをそのまま記述するだけ。
dictionary comprehensionとenumerateを使って座標と値のdictionaryを一行で作れるのが地味にうれしい。
第4問
どの5つをとって足しても合成数になるユニークなn個の素数のリストを作成する問題。
全部1で終わる数字なら5つとって足し合わせると確実に5の倍数になる。エラトステネスの篩を作って小さい順に1で終わる素数をn個取り出すだけ。
from itertools import count, islice def eratosthenes(): d = {} for q in count(2): p = d.pop(q, None) if p: x = next(x for x in count(q + p, p) if x not in d) d[x] = p else: yield q d[q*q] = q n = int(input()) ones = (prime for prime in eratosthenes() if prime % 10 == 1) n_ones = islice(ones, n) print(' '.join(map(str, n_ones)))
素数をつくる関数が定義されていればロジックは最後の三行でまとめられてしまう。
そういえば素数のリストを作成する関数が標準ライブラリに入っていないのはちょっと不思議だろうか?あまり汎用性がないという判断なのかな。他の言語でも標準ライブラリに入っているという話はあまり聞かないかも。