codeFlyer （bitFlyer Programmin...) C.徒歩圏内

Python Atcoder

問題文

N 個の都市があり、1,2,…,N の番号がついています。
これらの都市はこの順に一直線上に並んでいます。
各 i (1≤i≤N) について、都市 i の座標は Xi です。

高橋くんは都市 i と都市 j の間の移動手段を以下のように選びます。

  • 都市 i と都市 j の距離 |Xi−Xj| が D 以下であれば、徒歩で移動する。
  • そうでない場合、電車で移動する。

3 つの都市 (の番号) の組 (i,j,k) であって、以下の条件を満たすものの個数を求めてください。

  • i<j<k
  • 高橋くんは都市 i と都市 j の間、および都市 j と都市 k の間を徒歩で移動する。(Rule 1)
  • 高橋くんは都市 i と都市 k の間を電車で移動する。(Rule 2)

制約

  • 3≤N≤105
  • 0≤Xi≤109 (1≤i≤N)
  • Xi<Xi+1 (1≤i<N)
  • 0≤D≤109

考え方

for文を二重構造にしたらそれぞれの都市間を徒歩で移動するか電車で移動するかの情報を
配列で持てる...と思ったけど制限時間がアウト

3つの都市番号の組のうちjに着目して、
【jより左側にある距離D以内の都市の数】✖︎【jより右側にある距離D以内の都市の数】(=Sとする)を求め、 iに着目して、【iの右側にある距離D以内の2点の組み合わせの数】＝(Tとする)を求める。
すると、S-Tが答えになってくれる。

すなわちRule1を満たすものの個数を求めた後、Rule2を満たしていないものを引く。 二分探索で距離D以内の都市の配列番号を求めると良い。 ライブラリbisect...便利...

ソースコード

import bisect

N,D=map(int,input().split())
X=list(map(int,input().split()))

ans=0


#iは配列番号(=自分より左にある点の数)

for i in range(N-1):
 l=bisect.bisect_left(X,X[i]-D)     #lは自分より左にある点で距離D以内にない点の数と同義
 r=bisect.bisect_right(X,X[i]+D)     #rは下図参考
 ans+=(i-l)*(r-i-1)      #Rule1 左側の都市の数*右側の都市の数(rから自分と左側の都市の数を引く) 
 ans-=(r-i-1)*(r-i-2)//2  #Rule2 右側の都市の数から2個選ぶ組み合わせの数

print(ans)

具体例図

f:id:inaba00032:20180606033707j:plain

その他