🐛

[鉄則A13] Close Pairs (しゃくとり法)

2024/05/05に公開

Ruby

しゃくとり法^[1]を学ぶ基本問題。

問題を要約する

nC2 の組み合わせの総数は？
ただし二つの値の差が K 以下の組み合わせのみを有効とする

しゃくとり虫としゃくとり法

本物のしゃくとり虫とは少しイメージが異なる。虫は頭とお尻がくっつくような動きを繰り返すが、法は最後を除いてくっつかない。

例1でしゃくとり法を学ぶ

配列 A は昇順に並んでいる。

0	1	2	3	4	5	6
11	12	16	22	27	28	31

次に l = 0 r = 0 として、

l の場所にある 11 と r + 1 の場所にある値 12 を比べる
差が (12 - 11) <= K(10) であれば r を進めて再度 (1) から繰り返す
r が止まったら l を進める

これを A.size 回繰り返すと l r は、

l	r	個数 (r - l)
0	2	2
1	3	2
2	3	1
3	6	3
4	6	2
5	6	1
6	6	0

となり個数の合計 11 が答えになる。

全探索版 (TLE)

N, K = gets.split.collect(&:to_i)                 # => [7, 10]
A = gets.split.collect(&:to_i)                    # => [11, 12, 16, 22, 27, 28, 31]
p A.combination(2).count { |a, b| (b - a) <= K }  # => 11

組み合わせの差分が K 以下のものを言葉通りに実装したもの。

二分探索版 (AC)

128 ms

N, K = gets.split.collect(&:to_i)  # => [7, 10]
A = gets.split.collect(&:to_i)     # => [11, 12, 16, 22, 27, 28, 31]

def bsearch_from(left)
  l = left
  r = A.size - 1
  while l <= r
    m = (l + r) / 2
    diff = A[m] - A[left]
    if diff > K
      r = m - 1
    else
      l = m + 1
    end
  end
  [l, r]                           # => [3, 2], [4, 3], [4, 3], [7, 6], [7, 6], [7, 6], [7, 6]
  # 最終的に反転して l が境界の外側で r が境界の内側になるため l を返す
  r                                # => 2, 3, 3, 6, 6, 6, 6
end

sum = A.each_index.sum do |l|
  r = bsearch_from(l)              # => 2, 3, 3, 6, 6, 6, 6
  r - l                            # => 2, 2, 1, 3, 2, 1, 0
end
p sum                              # => 11

この方法でも一応解ける。

既存メソッドを使った二分探索 (AC)

117 ms

N, K = gets.split.collect(&:to_i)           # => [7, 10]
A = gets.split.collect(&:to_i)              # => [11, 12, 16, 22, 27, 28, 31]
sum = A.each.with_index.sum do |lv, l|
  r = A.bsearch_index { |e| (e - lv) > K }  # => 3, 4, 4, nil, nil, nil, nil
  r = (r || A.size) - 1                     # => 2, 3, 3, 6, 6, 6, 6
  r - l                                     # => 2, 2, 1, 3, 2, 1, 0
end
p sum                                       # => 11

「K より大きい値のインデックス」で求めた値は境界の外側なので -1 している。

解説のしゃくとり法 (AC)

78 ms

N, K = gets.split.collect(&:to_i)               # => [7, 10]
A = gets.split.collect(&:to_i)                  # => [11, 12, 16, 22, 27, 28, 31]
r = []
A.each.with_index do |lv, i|
  if i == 0
    # 前回の位置はわからないので左端から進める
    r[i] = i
  else
    # 前回の最後の位置から進める
    # i から r[i - 1] の範囲はすでに K 以下になるのはわかっているため、
    # r[i - 1] を貰って進めた方が効率がよい
    r[i] = r[i - 1]                             # => 2, 3, 3, 6, 6, 6
  end
  # 人差指を伸ばす
  while r[i] < A.size - 1
    diff = A[r[i] + 1] - lv                     # => 1, 5, 11, 10, 15, 11, 5, 6, 9
    if diff > K
      break
    end
    r[i] += 1                                   # => 1, 2, 3, 4, 5, 6
  end
  r                                             # => [2], [2, 3], [2, 3, 3], [2, 3, 3, 6], [2, 3, 3, 6, 6], [2, 3, 3, 6, 6, 6], [2, 3, 3, 6, 6, 6, 6]
end
r                                               # => [2, 3, 3, 6, 6, 6, 6]
counts = r.collect.with_index { |e, i| e - i }  # => [2, 2, 1, 3, 2, 1, 0]
p counts.sum                                    # => 11

r[i] は差が K 以下になるいちばん大きな値の位置を指している。次の位置に進めることができるのを確認してから +1 しているので r[i] は常に範囲の内側になっている。また二分探索よりも速い。

リファクタリング後 (AC)