A13 - Close Pairs

入出力

1
2
3
let N,K = stdin.ReadLine().Split() |> (fun x -> int x.[0], int64 x.[1])
let Ia = stdin.ReadLine().Split() |> Array.map int64
solve N K Ia |> stdout.WriteLine

方針

尺取り法で対処します. ここでは尺取り法は解説しません. オンライン上の資料としては例えばここを参考にしてください.

C/C++/Rustなど速い言語, またはHaskellでは問題ないようですが, AtCoder上のF#では単純な探索でTLEしてしまったため, 尺取り法の探索部分を二分探索で書く必要がありました.

2022/12時点の私の実装力だと, 二分探索ではまり倒したために力づくの部分があり, あまり綺麗なコードになっていません. いつかもう少しすっきり書き直したいです.

解説

前処理

今回は親切に入力の配列がソートされているため不要です. 実際には必要に応じてソートします.

大枠

配列の各添字に対して条件をみたす最長の添字を取ればよいため, 各iごとに最長の添字を探す関数をsearchとすれば実装の本体は次のように書けます.

1
  [|0..N-1|] |> Array.sumBy (fun i -> search i |> function | Some j -> (j-i |> int64) | None -> 0L)

条件をみたす添字がない可能性があるため, searchの返り値はOptionにしています. 特に条件をみたす添字jに対してj-iを積み, そうでない場合は0にして和を取れば求める組み合わせの総和が得られます. ここで一般に総和の値はintの範囲を越えるため, j-iにint64をかませる必要があります. 実際これでREをくらって原因がわからず30分ほどはまり倒しました.

何はともあれあとはsearchを実装すれば終わりです.

二分探索に入る前にまずは条件をみたす添字があるかどうかを判定します. ここでは次のように書きます.

1
  let search i = if i=N-1 || K<Ia.[i+1]-Ia.[i] then None else Some "二分探索"

i=N-1の場合は後続がありません. また既にソートされているため, 配列の次の添字との差が既にKを越えていれば条件をみたす添字がありません. したがってこの二者の場合はNoneを返します. あとは探索値が必ず存在する仮定のもとで二分探索します.

二分探索の実装

まず真偽計算用の関数を用意します.

1
  let p i j = Ia.[j] - Ia.[i] <= K

各iごとにこれを使って次のように計算します.

1
  let search i = if i=N-1 || K<Ia.[i+1]-Ia.[i] then None else Some (bsearch (p i) i (N-1))

あとはpまたはp iを使って次のように二分探索を書きます.

1
2
3
4
5
6
  let rec bsearch pi l r =
    if r<=l then if pi l then l else l-1
    else
      let m = (l+r)/2
      if pi m then bsearch pi (m+1) r
      else bsearch pi l m

いつも通り終了条件はr <= lです. 実際に実行してみるとわかるように, 上記の実装ではlが適切な添字の値になるとは限りません. そこでpi l = p i lを判定して条件をみたすならlを, みたさない場合は1を引いたl-1にします. (TODO: ここでスカっとlを返したい.)

次は再帰部分です. これもいつも通りまずは中点を取るべくm = (l+r)/2を取ります. もしIa.[m]が条件をみたすならm+1以上r以下から新たに添字を探します. もしIa.[m]が条件をみたさないからl以上m以下から新たに添字を探します.