Arantium Maestum

プログラミング、囲碁、読書の話題

AtCoder Beginner Contest 102に参加してみた

第1問

abc102.contest.atcoder.jp

abc102.contest.atcoder.jp

入力値nと2の最小公倍数を求める問題。nが偶数ならn、そうじゃないなら2*n

n = int(input())
print(2*n if n%2 else n)

第2問

abc102.contest.atcoder.jp

abc102.contest.atcoder.jp

Aの(添字の)異なる 2 要素の差の絶対値の最大値を求めてください。

と言われるとややこしく感じてしまうが「2要素の差が一番大きい」なのでmaxとminの差。

入力値に「Aの要素が2つ以上」という制約がある以上、要素の値がすべて同じだったとしても「添字の異なる2要素」は満たせる。

n = int(input())
xs = [int(x) for x in input().split()]
print(max(xs) - min(xs))

第3問

abc102.contest.atcoder.jp

abc102.contest.atcoder.jp

Ai と b+iの差

と考えるとちょっとややこしいかもしれないがabs(Ai - (b+i) = abs((Ai - i) - b)なのでまずAi-iの数列Bを作ってしまう。

その数列Bの各要素と任意の数bの距離の和を最小化したい。

たとえばbがBのどの要素よりも小さかったとする。そうするとb+1はBのすべての要素に1ずつ近くなるので明らかにbより勝る。

bがBの最小の要素と同じ値だったとする(Bの最小の要素が単一だと仮定する)。するとb+1はその最小の要素からの距離は1増え、n-1個の要素との距離は1ずつ減る。

この「1ずれるごとに距離が1増える要素と距離が1減る要素の数」がちょうどバランスするところが最善であり、そのポイントでの距離の和が答えとなる。

それは中央値を構成する要素が1つの場合はちょうどその要素の値となるし、もし中央値が二つの要素の平均だとしたら、その二つの要素の値の間ならどこでも最善になる。

なのでてっとりばやく中央値を使える:

from statistics import median
 
n = int(input())
xs = [int(x) for x in input().split()]
 
ys = [x-i-1 for i, x in enumerate(xs)]
m = int(median(ys))
 
print(sum(abs(y-m) for y in ys))

第4問

abc102.contest.atcoder.jp

abc102.contest.atcoder.jp

この問題はコンテスト中には解けなかった・・・ しゃくとり法や左右から端を狭めていく方法などをずっと考えていて最後に苦し紛れの実装でWAを出したりしながらタイムアップ。悔しいので解説は見ずに1日ぼーっと考えたりしていた。

累積和をとると便利なのは最初から考えていたのだが、「真ん中をまず割ってみる」というところに発想がいかなかったのが敗因。

真ん中を割ってしまえば、「左右の配列を個別にどう割るか」という二つの小さい問題になる。

まず左の配列について考えてみる:

  • 配列をどう割るのが最善か、と考えると、ここでも和の差を最小化するのが一番だとわかる
  • 和の差が最小化されるためには一つ一つの和ができるだけ平均に近ければいい
  • 累積和を平均値で二分探索すればいい
  • ただし二分探索で得られた添字の周りが答えの可能性もあるので-1, +1も調べて和の差が最小化されるポイントを選ぶ

右の配列も似た論理だ。ただし累積和がすこしややこしい(左の配列の総和を差し引く必要がある)。

あとは左右の配列の結果を組み合わせて最大と最小の和の差をとれば、あるポイントを真ん中にして割った場合の最小の和の差がわかる。

つまり真ん中が決定していれば、最小の和の差を2*O(log N/2)で計算できる。

あとは真ん中候補すべてに対してこの計算をして最小の結果が全体の答え。なので全体でO(N logN)。

from itertools import accumulate, islice
from bisect import bisect_left
 
n = int(input())
xs = [int(x) for x in input().split()]
 
ys = list(accumulate(xs))
zs, total = ys[:-1], ys[-1]
 
res = max(xs) - min(xs)
for i, z in islice(enumerate(zs), 1, n-1):
    j = bisect_left(zs, z//2)
    splits = ((zs[j+k], z - zs[j+k]) for k in (-1, 0, 1) if 0 <= j+k <= i)
    a, b = min(splits, key=lambda s:abs(s[0]-s[1]))
 
    j = bisect_left(zs, (total - z)//2 + z)
    splits = ((zs[j+k] - z, total - zs[j+k]) for k in (-1, 0, 1) if i <= j+k <= n-2)
    c, d = min(splits, key=lambda s:abs(s[0]-s[1]))
 
    minn, _, _, maxx = sorted((a,b,c,d))
    if maxx - minn < res:
        res = maxx - minn
 
print(res)

Pythonだとテストケースにかかる時間が最長で1800msとかなりギリギリだ。うっかりj = bisect_left(zs, z/2)などと浮動小数点との比較で二分探索してしまうと簡単にTLEになる。今度OCamlで実装してどう書けてどれくらい速度がでるか試してみたい。