読者です 読者をやめる 読者になる 読者になる

ゆらのふなびと

競プロ, Python, C++

Codeforces Round #341 (Div. 2) B - Wet Shark and Bishops

競プロ 斜め defaultdict

問題

1000*1000の正方形のチェス盤がある。この上にn個のビショップを置く。

ある2つのビショップについて、それらが同じ対角線上にあるとき、ビショップは「互いを攻撃している」と言う(AとBの間に他のビショップCがあっても、AとBが同じ対角線上にあれば「攻撃している」である)

n個のビショップと、それぞれの位置(x_i, y_i)が与えられるので、互いに攻撃しているビショップの組の数を答えよ。

制約

1 <= n <= 200000

解法

ある2つのビショップA(x1, y1), B(x1, y2)について、AとBが互いに攻撃している、つまり同じ対角線上にあるとは、以下と同値である。

x1 + y1 = x2 + y2 or x1 - y1 = x2 - y2

ここで、x1 + y1 = x2 + y2 = k を満たすビショップの個数が m 個だったとすると、このc個の中に攻撃しあっているビショップのペアは mC2 個である(m個のビショップから2つ選んでペアを作る場合の数だから) x1 - y1 = x2 - y2 = k のときも同様である。

よって全体としては以下のように処理すればよい。

  1. x, yを読み込みながら各x+y, x-yの個数をカウントしていく
  2. カウントした個数mについて、mC2を足し上げていく
from collections import defaultdict as ddict

def nC2(n):
    return n * (n - 1) // 2

SIZE = 1000

N = int(input())
plus = ddict(int)  #x+y
minus = ddict(int)  #x-y
for i in range(N):
    x, y = map(int, input().split())
    plus[x + y] += 1
    minus[x - y] += 1

ans = 0

for x in plus.values():
    ans += nC2(n = x)

for x in minus.values():
    ans += nC2(n = x)

print(ans)

x+y, x-yの個数のカウントにはdefaultdictを用いている。ただのdictだと数え上げをしたいとき、

if x in dic:
    dic[x] += 1
else:
    dic[x] = 1

のように書かなければいけないが、defaultdictではキーが存在しなかった場合自動で初期化してくれる。何で初期化するかは最初に与えることができる。たとえば上のコードではdefaultdict(int)により0で初期化するようになっているし、valueをリストにしたい場合はdefaultdict([])defaultdict(list)(訂正:2016/02/01 19:46)でOK。

Q.ボード全走査じゃだめ?

ここからはおまけ。

実は、私は本番では上のやり方ではなく、ボード上のすべての対角線を走査してカウントするという方法でACした。そのコードが以下。

def nC2(n):
    return n * (n - 1) // 2

SIZE = 1000

N = int(input())
a = [[0] * SIZE for i in range(SIZE)]
for i in range(N):
    x, y = map(int, input().split())
    a[x-1][y-1] = 1

ans = 0

#1
for sx in range(SIZE):
    t_cnt = 0
    
    x = sx; y = 0;
    while x >= 0:
        t_cnt += a[x][y]
        x -= 1; y += 1
        
    ans += nC2(t_cnt)

#2
for sy in range(1, SIZE):
    t_cnt = 0
    
    x = SIZE - 1; y = sy;
    while y < SIZE:
        t_cnt += a[x][y]
        x -= 1; y += 1
        
    ans += nC2(t_cnt)

#3
for sx in range(SIZE):
    t_cnt = 0
    
    x = sx; y = 0;
    while x < SIZE:
        t_cnt += a[x][y]
        x += 1; y += 1
        
    ans += nC2(t_cnt)

#4
for sy in range(1, SIZE):
    t_cnt = 0
    
    x = 0; y = sy;
    while y < SIZE:
        t_cnt += a[x][y]
        x += 1; y += 1
        
    ans += nC2(t_cnt)

print(ans)

まあとりあえず、長い。実行時間・メモリ使用量も先の解法の方が良いです(下図参照)

f:id:pakapa104:20160201142606p:plain

斜めの判定は x+y と x-y だけでOK。というのがこの問題で得たことでした。