概要

一直線上に並んだブースのうち、座標の差が \(D\) 以下であるすべてのペア \((i, j)\) について、集客力の積 \(S_i \times S_j\) の総和を効率的に求める問題です。

考察

素朴なアプローチとその問題点

最も単純な方法は、すべてのペア \((i, j)\)（\(i < j\)）を列挙し、\(|X_i - X_j| \leq D\) かを判定して積を足す方法です。しかし、ペアの数は最大で \(\frac{N(N-1)}{2} \approx 2 \times 10^{10}\) にもなり、TLE（制限時間超過）になります。

重要な気づき

ソートと二分探索を組み合わせることで高速化できます。

ブースを座標の昇順にソートすると、ブース \(i\) とコラボレーションが成立するブース \(j\)（\(j > i\)）は、\(X[j] \leq X[i] + D\) を満たす連続した範囲になります。この範囲の右端は二分探索で \(O(\log N)\) で求められます。

さらに、求めたい値は次のように変形できます：

\[\sum_{\text{valid } (i,j)} S_i \times S_j = \sum_{i} S_i \times \left( \sum_{j=i+1}^{r_i} S_j \right)\]

ここで \(r_i\) はブース \(i\) とコラボ可能な最も右のブースのインデックスです。内側の和（\(S_j\) の区間和）は累積和（prefix sum）を使えば \(O(1)\) で計算できます。

具体例

例えば、ソート後のブースが以下のとき（\(D = 3\)）：

インデックス	座標 \(X\)	集客力 \(S\)
0	1	2
1	3	5
2	4	3
3	8	1

• \(i=0\)：\(X[j] \leq 1+3=4\) なので \(j=1,2\) が対象 → \(S[0] \times (S[1]+S[2]) = 2 \times 8 = 16\)
• \(i=1\)：\(X[j] \leq 3+3=6\) なので \(j=2\) が対象 → \(S[1] \times S[2] = 5 \times 3 = 15\)
• \(i=2\)：\(X[j] \leq 4+3=7\) なので対象なし → \(0\)

総和は \(16 + 15 = 31\) となります。

アルゴリズム

1. ブースを座標 \(X_i\) の昇順にソートする。
2. 集客力 \(S\) の累積和配列 prefix_S を構築する。
3. 各ブース \(i\)（\(0 \leq i < N\)）について：
   • 二分探索で \(X[j] \leq X[i] + D\) を満たす最大の \(j\) を求める（bisect_right を使用）。
   • 累積和を用いて \(S[i+1] + S[i+2] + \cdots + S[r_i]\) を \(O(1)\) で取得する。
   • \(S[i] \times (\text{区間和})\) を答えに加算する。
4. 総和を出力する。

計算量

• 時間計算量: \(O(N \log N)\)
  • ソートに \(O(N \log N)\)、各ブースに対する二分探索に \(O(\log N)\)（合計 \(O(N \log N)\)）
• 空間計算量: \(O(N)\)
  • 座標・集客力の配列、累積和配列の格納に使用

実装のポイント

• ソート後は座標が昇順なので、\(|X_i - X_j| = X_j - X_i\)（\(j > i\)）と絶対値を外せる。

• bisect_right(X, X[i] + D) は「\(X[i] + D\) より大きい最初の位置」を返すので、\(-1\) して右端のインデックスを得る。

• 答えが非常に大きくなり得るが、Python は多倍長整数を標準でサポートしているためオーバーフローの心配は不要。

• 入力が大量になる可能性があるため、sys.stdin.buffer.read() で一括読み込みすると高速。

  ソースコード

import sys
from bisect import bisect_right

def main():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(input_data[idx]); idx += 1
    D = int(input_data[idx]); idx += 1
    
    booths = []
    for i in range(N):
        x = int(input_data[idx]); idx += 1
        s = int(input_data[idx]); idx += 1
        booths.append((x, s))
    
    # Sort by coordinate
    booths.sort()
    
    X = [b[0] for b in booths]
    S = [b[1] for b in booths]
    
    # For each booth i, find all booths j > i such that X[j] - X[i] <= D
    # Since sorted, X[j] >= X[i], so condition is X[j] <= X[i] + D
    
    # We need sum over all valid pairs (i,j) of S[i]*S[j]
    # = sum_i S[i] * (sum of S[j] for j in range where i < j and X[j] <= X[i]+D)
    
    # Use prefix sums of S
    prefix_S = [0] * (N + 1)
    for i in range(N):
        prefix_S[i + 1] = prefix_S[i] + S[i]
    
    total = 0
    for i in range(N):
        # Find rightmost j such that X[j] <= X[i] + D
        right = bisect_right(X, X[i] + D) - 1
        # Sum S[j] for j from i+1 to right
        if right > i:
            sum_s = prefix_S[right + 1] - prefix_S[i + 1]
            total += S[i] * sum_s
    
    print(total)

main()

---

この解説は claude4.6opus-thinking によって生成されました。