CodeIQ問題のHaskell解答

うまく通ったHaskellでの解答を載せてみます。（回答期限は終了しています）
最後の方にRuby版も書いてみました。

問題

つなぎ合わせたときの長さLと、N個(1≦N≦5000)の棒の長さが標準入力から与えられるとき、N個の棒の中から3つをつなぎ合わせて長さがLになる組み合わせの総数を求めるプログラムを書いてください。ただし、個々の棒の長さや、つなぎ合わせた長さ(L)は正の整数で、32bit整数で十分扱える範囲です。また、棒の長さはすべて異なるものとします。

【入力】
標準入力から、以下の形式で複数の整数値を読み込みます。

L
N
a1
a2
...

1行目のLは、つなぎ合わせた棒の長さです。
2行目のNは、棒の数です。
3行目以降のN行は個々の棒の長さで、長さはすべて異なるものとします。

【出力】
標準出力に、整数値を1つ出力してください。行末での改行は有無を問いません

プログラムの実行時間に制限があり、N=5000の時も1秒以内に計算終了する必要があります。

組み合わせ総数

愚直に全てのパターンを組み合わせた場合の母数は「N個の異なる値から3個を順序を考慮せず取り出す」組み合わせなので (5000×4999×4998)/(3×2×1) = 20820835000となり約200億通りになります。このまま走査しては当然間に合いません。

愚直な方法による解答例

とはいえ、まずはシンプルに順番に走査する方法で解いてみましょう。

• リストから値 x を一つ取る
• x より後方のリストから値 y を一つ取る
• y より後方のリストから、L-x-y に等しい値を探す
• 存在すれば 1, 存在しなければ 0 を返す

リストを前から順にペア (x,y) を選んでいけば総和が L に一致する組み合わせの数が出るはずです。

import Data.List

main = do
  l <- fmap toInt $ getLine -- 標準入力からLを読み込む
  n <- fmap toInt $ getLine -- 標準入力からNを読み込む(使わない
  bars <- fmap (fmap toInt . lines) $ getContents -- 残りのリストを取得
  print (sum $ fmap (findTripleSum l) $ tails bars) -- tails で部分リストを順番に調べる

-- Int化
toInt :: [Char] -> Int
toInt xs = read xs :: Int

-- リストの先頭をピックアップして、 和がL-xになる組み合わせを後方のリストから探す  
findTripleSum :: Int -> [Int] -> Int
findTripleSum _ [] = 0
findTripleSum _ (_:[]) = 0
findTripleSum i (x:xs) = sum $ fmap (findDoubleSum (i-x)) $ tails xs

-- リストの先頭をピックアップして、L-x-yに等しい物を探す
findDoubleSum :: Int -> [Int] -> Int
findDoubleSum _ [] = 0
findDoubleSum i (x:xs) = if elem (i-x) xs then 1 else 0

tails はリストの前方から一要素ずつdropした部分リストのリストを作るので、これを2回利用して2重ループを作っています。

Prelude> import Data.List
Prelude Data.List> tails [1,2,3]
[[1,2,3],[2,3],[3],[]]

当然のことながら、ほぼ全組み合わせを探すのでこのアルゴリズムでは遅いです。

両端からバランスを取りながら探す

以下提出した解答です。

入力される値は高々5000個なので、ソート自体は瞬間的に終わります。
ただ、ソートしたからといって前述のアルゴリズムでは遅いので、2重ループを作ってはいけません。

余談ですが箸の両端を両手の人差し指に乗せて、中央に向けて両指の間隔を縮めてみたことはありますか？
指にかかる摩擦力が重心に近い方が大きくなるため進まなくなり、自然とバランスを取ったまま箸を落とさずに重心まで行けるのです。

これと同じように、ソートされたリストに対して両端から同時に走査します。

import Data.List

main :: IO ()
main = do
    l <- fmap toInt $ getLine
    n <- fmap toInt $ getLine
    sortedBars <- fmap (sort . fmap toInt . lines) $ getContents -- ソートしておく
    print (countFixedLengthTripleBars l sortedBars)

toInt :: [Char] -> Int
toInt xs = read xs :: Int

-- tailsで部分リストを作って処理して、結果をsum
countFixedLengthTripleBars :: Int -> [Int] -> Int
countFixedLengthTripleBars length = sum . fmap (countTripleSum length) . tails

-- 先頭の値 x を使って、残りのリストを(L-x)/2を基準に分割して走査する
countTripleSum :: Int -> [Int] -> Int
countTripleSum _ [] = 0
countTripleSum l (x:xs) = countDoubleSum 0 s (filter (< h) xs) (reverse $ filter (>= h) xs) -- 大きい方のリストは reverse しておく
    where s = l - x
          h = ceiling ((fromIntegral s)/2)

-- s (= L-x) に一致するペアを大小のリストを走査しながら探す
-- countで見つけた数を更新して末尾再帰にする
countDoubleSum :: Int -> Int -> [Int] -> [Int] -> Int
countDoubleSum count _ [] _ = count
countDoubleSum count _ _ [] = count
countDoubleSum count s (low:lows) (high:highs)
  | sum == s = countDoubleSum (count+1) s lows highs
  | sum < s  = countDoubleSum count s lows (high:highs)
  | sum > s  = countDoubleSum count s (low:lows) highs
    where sum = low + high

リストの要素 x をピックアップして、残りのリストから両端の和を計算してみてバランスを取りながら進みます。

• L-x に等しければ +1して、両端のカーソルを1つずつ進める
• L-x より小さければ、小さい方のリストのカーソルをひとつ進める（和が大きくなる）
• L-x より大きければ、大きい方のリストのカーソルをひとつ進める（和が小さくなる）

実に手続き的なアルゴリズムですが、これにより一つのループがなくなって高速化されます。

Ruby版

一応Rubyでも書いてみました。こちらも1秒を切ります。
こっちのほうが読みやすいかな？

inputs = STDIN.read.lines.map{|l| l.chomp.to_i }

L = inputs.shift
N = inputs.shift
D = inputs.sort  # ソートしておく

count = 0

(0..(N-2)).each do |i|

  remain = L - D[i]

  # カーソルの初期位置
  low_cursor = i + 1
  high_cursor = N - 1

  # カーソルが中央で合流するまで
  while low_cursor < high_cursor do
    s = D[low_cursor] + D[high_cursor]
    if remain == s
      count += 1
      low_cursor += 1
      high_cursor -= 1
    elsif remain > s
      low_cursor += 1
    else
      high_cursor -= 1
    end
  end
end

puts count