yukicoder No.202 1円玉投げ

こういう問題は見た瞬間にバケット法を思いつけるようにならないとなー

問題はこちら。

問題

机の上にN(≦ 10⁵)枚の1円玉を投げていく。1円玉は半径10で、立つことはない。i番目の1円玉を投げたところ、中心の座標がx_i, y_i(ともに0以上20000以下）となったとする。このときすでに存在している1円玉と重なってしまった場合は、この投げた1円玉は取り除く。N枚投げ終わったあと、机の上に残っているコインの枚数を答えよ。

解法

机の上に残っているすべてのコインとの重複を調べると最悪O(N²)かかってしまう。またどの座標にコインがすでにあるかを記録しようとしてもO(20000²)のメモリが必要になってしまいMLEとなってしまう。

この問題はバケット法を使えば解くことができる。

たとえば平面を１辺50のブロックに分けて、机の上に残った1円玉がどのブロックに属すかを記録していく。 こうすればメモリの問題も解消でき、重複を調べるべき1円玉の枚数も減らすことができる。
重複判定は、投げられた1円玉が属するブロックとその周りの9ブロックにある1円玉について調べればいい。
1辺50だと1ブロックに最大10個くらいしか1円玉は存在しえないので、合計10ブロックを調べるのに最大でも100回の演算で済む。

size = 50
def dist(ax, ay, bx, by):
    return (ax - bx) ** 2 + (ay - by) ** 2

def check(x, y):
    for dy in xrange(-1, 2):
        for dx in xrange(-1, 2):
            bx, by = x / size + dx, y / size + dy
            if 0 <= bx < 20000 / size + 1 and 0 <= by < 20000 / size + 1:
                for cx, cy in block[by][bx]:
                    if dist(x, y, cx, cy) < 400:
                        return False
    return True
    
N = int(raw_input())
block = [[[] for i in xrange(20000 / size + 1)] for j in xrange(20000 / size + 1)]

ans = 0
for loop in xrange(N):
    x, y = map(int, raw_input().split())
    if check(x, y):
        ans += 1
        block[y / size][x / size].append((x, y))
print ans

まとめ

典型なのに何度もWAを出してしまった