A05 - Three Cards

入出力

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

方針

これもよく出てくる処理です. 単純に三重ループを回すとTLEで時間内に終わりません. 時間内に終わらせる工夫が必要です.

解説

単純な処理の累積

三つ目のループをどうにかして回避する必要があります. 発想を転換して二つ目までのループで次のような和を作りましょう.

1
  [| for i in 1..N do for j in 1..N do yield K-i-j |]

この配列はN^2個の要素を持っています. 仮に三つ目のループを回すとしてその変数をxとすると, x = K-i-jが所定の範囲である1からNにおさまっていれば条件をみたします. そこでfilterを使ってチェックしましょう.

1
2
  [| for i in 1..N do for j in 1..N do yield K-i-j |]
  |> Array.filter (fun x -> 1<=x && x<=N)

これで条件をみたす要素だけが抜き出せました. あとはこの配列の長さを取れば終わりです.

1
2
3
  [| for i in 1..N do for j in 1..N do yield K-i-j |]
  |> Array.filter (fun x -> 1<=x && x<=N)
  |> Array.length

forの中でフィルターする

他の言語と同じくforの内包表記の中で次のようにフィルターできます.

1
  [| for i in 1..N do for j in 1..N do let x=K-i-j in if 1<=x && x<=N then x |]

あとは同じく配列の長さを取れば終わりです.

1
2
let solve N K =
  [| for i in 1..N do for j in 1..N do let x=K-i-j in if 1<=x && x<=N then x |] |> Array.length

filterでN^2個の配列を再びチェックしなくていい分, こちらの方が二割ほど速いです.