Arantium Maestum

失敗１：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では非常に深い再帰の問題は再帰リミットだけではなく、関数スタックが大きいのでメモリ使用が厳しいという点もある。末尾呼び出し最適化もないし。

arc031.contest.atcoder.jp

というわけで、まあ再帰しなければいいよね。という単純な解決。

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 - バウムテスト

arc037.contest.atcoder.jp

無指向グラフにいくつ非巡回グラフが含まれているかを数える問題。

問題文をみると???だが実際に解法は非常に似通っている：

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として乗せることで対応している。

2018-06-18

蟻本初級編攻略 - 2-1 部分和

AtCoder Python 蟻本

定期的に「蟻本を全部読むぞ！」と思って読み始めては放置、を繰り返している。

そもそもAtCoderやCodeForceに参加することなく「ある程度読んでから競プロ試してみよう」と思っていたのが間違いだったのではないか。

ABC100にリアルタイム参加もしたし、今度こそ蟻本を攻略したい。

幸いこういう素晴らしい記事もある：

qiita.com

本を読み、AtCoderの問題を解き、その解法をブログに書いていこうと思う。

部分和問題

蟻本の冒頭あたりについては今年の三月にも記事を書いていた。

というわけで書いたブログとコードを引っ張り出してみる：

zehnpaard.hatenablog.com

が、なんだこれは・・・　ひどいねコレは。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の類題を解いてみる。

arc061.contest.atcoder.jp

これは枝刈りなしの全列挙なので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

abc079.contest.atcoder.jp

ロジックはまったく同じ：

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で最初の要素を取り出している。

2018-06-17

AtCoder Beginner Contest 095に事後参加してみた

AtCoder Python

この回はD問題がけっこう面白かった。ちょっとコードも解説も多めかもしれない・・・

第１問

700円のラーメンに味玉、チャーシュー、ネギのトッピング（各100円）するか否かを示す'o'か'x'かが三つ並んだ文字列を与えられて、必要な金額を計算する問題。

print(700 + 100*input().count('o'))

そのまんま。

第２問

ドーナッツの種類が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の和の量の粉で作って、残った粉でもっともコストの低いドーナッツを最大限作った時の数を返している。

第３問

二種類のピザ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)

第４問

円形のカウンターに置いてある寿司に、スタート地点からの時計回りでの距離（＝到達までのカロリー消費量）と寿司を食べた時に得られるカロリーの二つの値がある。時計回り、反時計回りに動き回って得られるカロリー値を最大化する問題。

時計回り、反時計回りの計算をなるべく対称にするため、「スタート地点からの時計回りでの距離のリスト」を「点と点の間の距離のリスト」に変換しておく。

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でその最大値をとっている。

方向転換も可能な場合

上記を踏まえて、方向転換について考えてみる。

少し考えると、二回以上方向転換をする意味がないことがわかる。

二回方向転換しても二度目の方向転換の後は「すでに寿司をとってしまった距離を再度歩いている」だけなので、「時計回り・反時計回り・時計回り」といった進みかたをするより「反時計回り・時計回り」と進んだほうが同じ寿司をより少ない距離でとれる。

なので方向転換を考慮に入れなければいけないのは一回のみ。

さらに「方向転換しない」というのは「まず一方向に０距離進んでから逆方向に転換する」という「一回方向転換」の特殊ケースとして扱える。

ということでこの問題は「まず時計回りにすすんでから方向転換して反時計回り」「まず反時計回りにすすんでから方向転換して時計回り」の二つのケースを調べるだけで済む。

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を使って効率良く記述できるのは嬉しい。

感想

コードを書くよりブログを書くほうが何倍も難しかった・・・　言語化って大変だな。

2018-06-17

AtCoder Beginner Contest 100に参加してみた

AtCoder Python

初リアルタイム参加。言語はPython3。

「ABCって二週間おきにあるのかな？次は6/23か」とのほほんとしていたら

今日は記念すべき100回目のBeginner Contest、AtCoder Beginner Contest 100なので、ぜひぜひ参加してみてくださいな！
Hello Worldが出力出来れば参加してみておっけーなレベルだよ！！！https://t.co/bCiCjMD2hT
— chokudai(高橋直大) (@chokudai) 2018年6月16日

こんなツイートが出ていて地味に驚く。

あと一週間あることを前提に過去問を最近のやつから順に解いていたのだが、最近出題されたものばっかり見てたということでなんとなく不利な気がしてしまう。（最近出題した問題と類似したものはすぐには出しにくいのでは？という考えから）

まあ仕方ないから当たって砕けろ、で初参加。

とりあえず３０分弱でミスなし全完で順位は86位。実力からすると完全にできすぎでビギナーズラック以外の何物でもない・・・

めでたく初回ランク上昇上限の400点をもらい茶色コーダーに。

第１問

ケーキを１６等分に切って、同じ人が隣り合わせのピースを取れないとして、二人が申告した数のピースを食べられるかどうかを判定する。

16等分と偶数で、ケーキをとる人数も２なので非常に楽。隣り合わせのピースはとれないのだから一人当たり最大で８しかとれず、お互いにとった数は干渉し合わない。

a, b = map(int, input().split())
print('Yay!' if a < 9 and b < 9 else ':(')

第２問

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するか/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関数はもうちょっときれいに書けそう。

第４問

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 プロトコルをうまく組み合わせるとけっこうきれいに、宣言的に記述できるように思う。

すごく楽しかった。次回もまた参加したい。

2018-06-16

AtCoder Beginner Contest 096に事後参加してみた

AtCoder Python

今日も今日とてAtCoder ABCを解く。

第１問

1月1日、2月2日、10月10日といったようにX月X日のようになっている日が年始からa月b日までに何日あったかを数える問題。

a, b = map(int, input().split())
print(a-(a>b))

ブール値が1と0に評価されることを利用している。各月に一回月と日が一致する日があり、その日を過ぎているかどうかをa > bで判定している。

第２問

a, b, cの整数のいずれかを二倍に変えるという操作をk回行い、和を最大化する問題。

a, b, c = sorted(map(int, input().split()))
k = int(input())
print(a+b+(c*(2**k)))

当然一番でかい数字をk回二倍にして足し合わせるのが最大。

第３問

隣接した二つのマスを黒く塗りつぶすという操作のみで、特定の白黒のパターンを再現できるか判定する問題。

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を一行で作れるのが地味にうれしい。

第４問

どの５つをとって足しても合成数になるユニークなn個の素数のリストを作成する問題。

全部１で終わる数字なら５つとって足し合わせると確実に5の倍数になる。エラトステネスの篩を作って小さい順に１で終わる素数を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)))

素数をつくる関数が定義されていればロジックは最後の三行でまとめられてしまう。

そういえば素数のリストを作成する関数が標準ライブラリに入っていないのはちょっと不思議だろうか？あまり汎用性がないという判断なのかな。他の言語でも標準ライブラリに入っているという話はあまり聞かないかも。

2018-06-15

AtCoder Beginner Contest 097に事後参加してみた

AtCoder Python

とりあえずこれから何週間か、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木など知ってはいたけれどもあまり引っ張り出して使ったことがない機能をいろいろ組み合わせたりできるのはかなり楽しい。これらを自然と使えるようになればコンテストの時間制限の中でもけっこう効率的できれいな解法に行き着きやすいのではないかと淡い期待を抱いている。

2018-06-14

AtCoder Beginner Contest 098に事後参加してみた

AtCoder Python

「事後参加」って変な言い回しだな。「の問題を解いてみた」が正しい・・・

例によってPythonで。

第１問

入力値二つの和、差、積のうちの最大値を出力する問題。本当に問題文のまま実装・・・

第２問

文字列sを分割して、「どちらにも出てくる文字」の種類数を最大化する問題。やはり問題文のまま実装・・・

というのはつまり文字列を分割した場合の二つ組を列挙して、文字をset化して、intersectionをとって、その長さの最大値を計算する、という流れ。

第３問

西向きの人、東向きの人が混ざった一列のなかからリーダーを選んで、他の人にリーダーの方を向いてもらう場合、振り向く必要のある人数を最小化する問題。

a = accumulate(c=='W' for c in 'X'+s[:-1])
b = reversed(list(accumulate(c=='E' for c in 'X'+s[-1:0:-1])))
print(min(x + y for x, y in zip(a, b)))

itertools.accumulateを使ってちょっとだけ気の利いたことをしている。

accumulateは例えば整数のリストをとって、各添字までの合計値のリストを返すような関数。（実際には引数・戻り値ともにリストではなくiterable/iteratorだが）。

各ポジションで東側にいる西向きの人数、西側にいる東向きの人数を足したシーケンスの最小値を探している。特に東向きのほうはreverseして計算して再度reverseしているのですこしコードがごちゃついているのが不満。iteratorはreversedできないのでいったんリストに直しているのも個人的には減点・・・

第４問