Haskell solution for Median Sort interactive problem from the Qualification Round of Google Code Jam 2021

MedianSort.hs

{-# LANGUAGE DeriveFoldable #-}

import System.IO
import Control.Monad
import Data.Foldable

data Tree a = Leaf | Node (Tree a) a (Tree a) deriving (Foldable)

singleton :: a -> Tree a
singleton a = Node Leaf a Leaf

fromThree :: (a, a, a) -> Tree a
fromThree (a, b, c) = Node (singleton a) b (singleton c)

type Input = (Int, Int)
type Output = Maybe [Int]

ask :: [Int] -> IO Int
ask ls = do
  putStrLn (unwords (map show ls))
  readLn

data From a = L a | C | R a

insert :: From Int -> Int -> Tree Int -> IO (Tree Int)
insert _ n Leaf = pure (singleton n)
insert _ n (Node (Node ll lx lr) x r) = do
  answer <- ask [lx, x, n]
  case () of
    _ | answer == lx -> do
          ll' <- (insert (R lx) n) ll
          pure (Node (Node ll' lx lr) x r)
      | answer == x -> do
          r' <- insert (L x) n r
          pure (Node (Node ll lx lr) x r')
      | answer == n -> do
          lr' <- insert (L lx) n lr
          pure (Node (Node ll lx lr') x r)
insert _ n (Node l x (Node rl rx rr)) = do
  answer <- ask [n, x, rx]
  case () of
    _ | answer == n -> do
          rl' <- insert (R rx) n rl
          pure (Node l x (Node rl' rx rr))
      | answer == x -> do
          l' <- insert (R x) n l
          pure (Node l' x (Node rl rx rr))
      | answer == rx -> do
          rr' <- insert (L rx) n rr
          pure (Node l x (Node rl rx rr'))
insert (L l) n (Node Leaf x Leaf) = do
  answer <- ask [l, x, n]
  case () of
    _ | answer == x -> do
          pure (Node Leaf x (singleton n))
      | answer == n -> do
          pure (Node (singleton n) x Leaf)
insert (R r) n (Node Leaf x Leaf) = do
  answer <- ask [n, x, r]
  case () of
    _ | answer == n -> do
          pure (Node Leaf x (singleton n))
      | answer == x -> do
          pure (Node (singleton n) x Leaf)


solve :: Int -> IO (Tree Int)
solve 3 = do
  x <- ask [1..3]
  pure . fromThree $ case x of
    1 -> (2, 1, 3)
    2 -> (1, 2, 3)
    3 -> (1, 3, 2)
solve n = do
  rest <- solve (n - 1)
  insert C n rest

main :: IO ()
main = do
  hSetBuffering stdin NoBuffering
  hSetBuffering stdout NoBuffering
  line <- getLine
  let [t, n, q] = map read (words line)
  replicateM_ t (do
    solution <- solve n
    response <- ask (toList solution)
    when (response == -1) (error "Wrong answer")
    )