AtCoderでScipyが使えると(頭が悪くても)役立つ話

昨日AtCoder Beginner Contest 223のC問題のいい解答が思い浮かばなかったので、Scipyで強引にやったら解けた。その話をまとめる。

AtCoder Beginner Contest (ABC) では、さまざまな言語を自分で選べる。そのため、言語によってはあらかじめ用意されているライブラリが使えるかどうかで有利不利が生まれることがある。本来はアルゴリズムとデータ構造の能力が問われる競技プログラミングだが、ライブラリの有無でCheatingできるパターンがある。

当該問題は、導火線の左右から火をつけたらどこの位置で火が出会うかを求めるという問題だった。模範解答では、導火線の片方から火をつけてすべて燃えきる時間の半分が、導火線の左右から火をつけたときに火が出会う時刻であることを利用して解いていた。この事実に気づけなかったので、延々と悩んでいた。

ところで、この問題は数値計算すると簡単に解けるなということも初めから気づいていた。下の図のように、y=f1(x)とy=f2(x)の交点を求めればいいので、f1(x)とf2(x)が求まればあとはソルバで解けばよい。普段だったら絶対こうやって解く。

頭の中で描いた解のイメージ

しかし、今回は次の二つの理由からこうやって解くのを躊躇していた。

  • ソルバであるscipy.optimize.rootと、線形補間をしてくれるscipy.interpolate.interp1dをAtCoderで使えるかどうか知らなかったこと
  • もし使えたとしても、時間計算量的に間に合うか想像がつかなかったこと
  • ちょっとズルっぽいこと

しかし時間が迫ってきたので、とりあえずやってみたら解けた。具体的なコードは次のとおり。interp1dで左右からの火の位置の関数を作り、それをソルバで解く。時間は全然間に合って拍子抜けした。

from itertools import accumulate
from scipy.interpolate import interp1d
from scipy.optimize import root
n = int(input())
abn = [[int(i) for i in input().split()] for _ in range(n)]
 
an = [a for a, b in abn]
tn = [a/b for a, b in abn]
tn2 = [a/b for a, b in abn[::-1]]
length = sum(an)
 
times1 = list(accumulate(tn))
times2 = list(accumulate(tn2))
 
len1 = list(accumulate(an))
len2 = list(accumulate(an[::-1]))
len2 = [length-_ for _ in len2]
 
times = sorted(times1+times2)
 
a1func = interp1d([0]+times1, [0] + len1)
a2func = interp1d([0]+times2, [length] + len2)
 
 
def fun(x):
    return a1func(x)-a2func(x)
 
 
sol = root(fun, [1])
[x] = sol.x
print(a1func(x))

時間がなくていい解も思いつかないとき、もしくはScipyのいいライブラリで問題が簡単に解けそうなときは試してみるとよいかもしれない。

コメント