Improvement over naive.hs

warrior.hs

import Control.Monad.Random
import Data.List
import System.Environment
import System.Exit
import System.Random
import System.IO

data Feedback = Smaller | Bigger | Done deriving (Show, Eq)
type Guess = Int
type Range = (Int, Int)

range = (1, 100)
weightedRange = [
    (1,   0.019609821304715197), (2,   0.014954327822406068), (3,   0.011497912909693942),
    (4,   0.011229933036558766), (5,   0.01029883434009694), (6,   0.008041497996981816),
    (7,   0.01029883434009694), (8,   0.01029883434009694), (9,   0.00917016616853938),
    (10,  0.01029883434009694), (11,  0.01029883434009694), (12,  0.006912829825424253),
    (13,  0.01029883434009694), (14,  0.01029883434009694), (15,  0.00917016616853938),
    (16,  0.01029883434009694), (17,  0.01029883434009694), (18,  0.008041497996981816),
    (19,  0.01029883434009694), (20,  0.01029883434009694), (21,  0.00917016616853938),
    (22,  0.01029883434009694), (23,  0.01029883434009694), (24,  0.01029883434009694),
    (25,  0.005784161653866689), (26,  0.01029883434009694), (27,  0.01029883434009694),
    (28,  0.00917016616853938), (29,  0.01029883434009694), (30,  0.01029883434009694),
    (31,  0.008041497996981816), (32,  0.01029883434009694), (33,  0.01029883434009694),
    (34,  0.00917016616853938), (35,  0.01029883434009694), (36,  0.01029883434009694),
    (37,  0.008605832082760596), (38,  0.008605832082760596), (39,  0.01029883434009694),
    (40,  0.00973450025431816), (41,  0.00973450025431816), (42,  0.01029883434009694),
    (43,  0.00917016616853938), (44,  0.00917016616853938), (45,  0.01029883434009694),
    (46,  0.00973450025431816), (47,  0.00973450025431816), (48,  0.01029883434009694),
    (49,  0.01029883434009694), (50,  0.007477163911203034), (51,  0.007477163911203034),
    (52,  0.01029883434009694), (53,  0.00973450025431816), (54,  0.00973450025431816),
    (55,  0.01029883434009694), (56,  0.00917016616853938), (57,  0.00917016616853938),
    (58,  0.01029883434009694), (59,  0.00973450025431816), (60,  0.00973450025431816),
    (61,  0.01029883434009694), (62,  0.008605832082760596), (63,  0.008605832082760596),
    (64,  0.01029883434009694), (65,  0.00973450025431816), (66,  0.00973450025431816),
    (67,  0.01029883434009694), (68,  0.00917016616853938), (69,  0.00917016616853938),
    (70,  0.01029883434009694), (71,  0.00973450025431816), (72,  0.00973450025431816),
    (73,  0.01029883434009694), (74,  0.01029883434009694), (75,  0.008041497996981816),
    (76,  0.008041497996981816), (77,  0.01029883434009694), (78,  0.00973450025431816),
    (79,  0.00973450025431816), (80,  0.01029883434009694), (81,  0.00917016616853938),
    (82,  0.00917016616853938), (83,  0.01029883434009694), (84,  0.00973450025431816),
    (85,  0.00973450025431816), (86,  0.01029883434009694), (87,  0.01029883434009694),
    (88,  0.006912829825424253), (89,  0.01029883434009694), (90,  0.01029883434009694),
    (91,  0.00917016616853938), (92,  0.01029883434009694), (93,  0.01029883434009694),
    (94,  0.008041497996981816), (95,  0.01029883434009694), (96,  0.01029883434009694),
    (97,  0.010101264865001205), (98,  0.012626581081251505), (99,  0.014954327822406068),
    (100, 0.019609821304715197)]

main :: IO ()
main = do
    args <- getArgs
    hSetBuffering stdout LineBuffering
    case args of
        ["pick"] -> pick
        ["guess"] -> guess
        _ -> usage

usage :: IO ()
usage = do
    putStrLn "usage: ( pick | guess )"
    exitWith (ExitFailure 1)

pick :: IO ()
pick = chooseNumber >>= putStrLn . show

chooseNumber :: IO Int
chooseNumber = evalRandIO (fromList weightedRange)

guess :: IO ()
guess = do
  subset <- chooseSubset 31
  interact $ formatGuesses . guess' subset . parseFeedback

parseFeedback :: String -> [Feedback]
parseFeedback = map parse . lines
    where parse l = case l of
            "+" -> Bigger
            "-" -> Smaller
            _  -> Done

formatGuesses :: [Guess] -> String
formatGuesses = unlines . map show

chooseSubset :: Int -> IO [Int]
chooseSubset n = do
    chosen <- evalRandIO $ chooseSubset' n weightedRange
    return $ sort chosen

chooseSubset' :: MonadRandom m => Eq a => Int -> [(a, Rational)] -> m [a]
chooseSubset' 0 _ = return []
chooseSubset' n options = do
    elem <- fromList options
    let remainingOptions = filter (\(option, _) -> option /= elem) options
    nextElems <- chooseSubset' (pred n) remainingOptions
    return $ elem : nextElems

guess' :: [Int] -> [Feedback] -> [Guess]
guess' [] _ = []
guess' candidates feedback = pivot : nextGuesses
    where pivot = candidates !! ((length candidates) `div` 2)
          nextGuesses = case feedback of
                (Bigger : nextFeedback) ->
                   guess' (filter (>pivot) candidates) nextFeedback
                (Smaller : nextFeedback) ->
                   guess' (filter (<pivot) candidates) nextFeedback
                _ -> []