Zenn
🔍

[鉄則A14] Four Boxes (半分全列挙)

2024/05/09に公開
Ruby
AtCoder
競プロ
二分探索
半分全列挙
tech

https://atcoder.jp/contests/tessoku-book/tasks/tessoku_book_n

半分全列挙の基本問題。

問題を要約する

  • 4つの箱にそれぞれ N 枚のカードが入っている
  • それぞれの箱から適当に1枚引く
  • 4枚のカードの合計が K になる可能性はあるか？

例1で問題を理解する

A B C D
3 4 10 1
9 7 20 2
17 9 30 3

から 9, 9, 30, 2 を引けば K(50) になる。

解き方

  1. A vs B を総当たりして足した P を作る
  2. C vs D を総当たりして足した Q を作る
  3. P + Q = K になるか調べる

最後は Q = K - P に変形して Q の中に K - P があるかを探す。

例1のシミュレーション

A B C D
3 4 10 1
9 7 20 2
17 9 30 3

から、

総当たり足し算の結果
P 7 10 12 13 16 18 21 24 26
Q 11 12 13 21 22 23 31 32 33

を作る。

続いて K(50) - 7 の 43 が Q に含まれるか探す。線形探索だと遅いので二分探索を行う。しかし 43 はなかったので次に進む。これを繰り返すと 18 のときに 50 - 18 の 32 が Q に見つかる。

全探索 (TLE)

N, K = gets.split.collect(&:to_i)                            # => [3, 50]
A, B, C, D = 4.times.collect { gets.split.collect(&:to_i) }  # => [[3, 9, 17], [4, 7, 9], [10, 20, 30], [1, 2, 3]]
ans = N.times.any? do |i|
  N.times.any? do |j|
    N.times.any? do |k|
      N.times.any? do |l|
        A[i] + B[j] + C[k] + D[l] == K
      end
    end
  end
end
ans                                                          # => true
puts ans ? "Yes" : "No"

最悪 O(N^4) かかってしまう。

解説の二分探索版 (AC)

444 ms
N, K = gets.split.collect(&:to_i)                            # => [3, 50]
A, B, C, D = 4.times.collect { gets.split.collect(&:to_i) }  # => [[3, 9, 17], [4, 7, 9], [10, 20, 30], [1, 2, 3]]
P = A.product(B).collect(&:sum)                              # => [7, 10, 12, 13, 16, 18, 21, 24, 26]
Q = C.product(D).collect(&:sum).sort                         # => [11, 12, 13, 21, 22, 23, 31, 32, 33]
ans = P.any? { |e| Q.bsearch { |q| (K - e) - q } }           # => true
puts ans ? "Yes" : "No"

find を bsearch で置き換えるなら find-any モードが適している。

集合 (AC)

473 ms
N, K = gets.split.collect(&:to_i)                            # => [3, 50]
A, B, C, D = 4.times.collect { gets.split.collect(&:to_i) }  # => [[3, 9, 17], [4, 7, 9], [10, 20, 30], [1, 2, 3]]
P = A.product(B).collect(&:sum)                              # => [7, 10, 12, 13, 16, 18, 21, 24, 26]
Q = C.product(D).collect(&:sum).to_set                       # => #<Set: {11, 12, 13, 21, 22, 23, 31, 32, 33}>
ans = P.any? { |e| Q.include?(K - e) }                       # => true
puts ans ? "Yes" : "No"

想定の解き方とは異なるが、集合にしてもほぼ同じ時間で解ける。Q は重複していても問題ない。

参考

Discussion