しゃくとり法（Close Pairs）

問題

$ N $ 個の整数があります。

整数は小さい順に $ A_1,A_2,…,A_N $ と並んでいます。

異なる２つの整数のペアを選ぶ方法は全部で $ N(N-1)/2 $ 通りありますが、その中で差が $ K $ 以下であるような選び方は何通りあるでしょうか。

[制約]
🔹$ 1 \leqq N \leqq 100000 $
🔹$ 1 \leqq K \leqq 10^9 $
🔹$ 1 \leqq A_1 < A_2 < … < A_N \leqq 10^9 $

解き方・ソースコード

$ A_t \leqq A_i + K $ を満たす最大の $ t $ を $ R_i $ とし、$ R_i $ の値を次のように計算します。

🔹手順１　$ i = 1 $ であれば $ R_i = 1 $ から、そうでなければ $ R_i = R_{i-1} $ からスタートする。
🔹手順２　差分 $ A_{R_{i}} - A_i $ が $ K $ を超えないギリギリまで、$ R_i $ を１増やす操作を繰り返す。

求めるペアの個数は $ (R_i - i) $ 通りで、答えは次式で計算できます。

$$ (R_1 - 1) + (R_2 - 2) + … + (R_{N-1} - (N -1)) $$

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
#--------- 入力例 -----------
N = 7       # 整数の数
K = 10      # 差をK以下にする
A = [9, 10, 14, 20, 25, 26, 29]     # 小さい順に並んだ整数
#----------------------------
# 配列の準備（R は 0 番目から始まる）
R = [ None ] * N

# しゃくとり法
for i in range(0, N - 1):
    # スタート地点を決める
    if i == 0:
        R[i] = 0
    else:
        R[i] = R[i - 1]

    # ギリギリまで増やしていく
    while R[i] < N - 1 and A[R[i] + 1]-A[i] <= K:
        R[i] += 1
    print(R)

# 出力
Answer = 0
for i in range(0, N - 1):
    Answer += (R[i] - i)
print('解：', Answer)

このアルゴリズムは、しゃくとり虫の動きにたとえてしゃくとり法と呼ばれています。

$ R_i $ を増やす操作は合計 $ N - 1 $ しか行われないので、計算量は $ O(N) $ となります。

[実行結果]

[2, None, None, None, None, None, None]

[2, 3, None, None, None, None, None]

[2, 3, 3, None, None, None, None]

[2, 3, 3, 6, None, None, None]

[2, 3, 3, 6, 6, None, None]

[2, 3, 3, 6, 6, 6, None]

解： 11

整数のペアの差が10以下になる組み合わせは、11通りあることを導き出すことができました。