Solution to ACM ECNA Regional, 2000, Problem B: Poly-polygonal Numbers

> import ReaderMonad
> import List
> import Monad

A bit of the mathematics of k-gonal numbers first.  Let N k n
be the n-th k-gonal number, with 1 being the 1st.

Triangle numbers: 1, 1+2, 1+2+3, 1+2+3+4, ...

N 3 1 = 1
N 3 (n+1) = N 3 n + (n+1)

Square numbers: 1, 1+3, 1+3+5, 1+3+5+7, ...

N 4 1 = 1
N 4 (n+1) = N 4 n + (2n+1)

In general,

N k 1 = 1
N k (n+1) = N k n + (k-2)n + 1

this is because the (n+1)th polygon comes from the nth polygon
by adding (k-2) sides of length n each.  This requires (k-2)n + 1
additional dots.

We compute the list of all k-gonal numbers in increasing order as
follows.  Consider the triangle numbers.  The list should be:

[1, 1+2, 1+2+3, 1+2+3+4, ...]

This is conveniently

scanl1 (+) [1,2,3,4,...]

Then we drop those below the given starting point.  We could instead
use algebra to determine the first k-gonal number above the starting
point with O(1) computation, but here we don't bother.

We will merge several such lists together.  In order to remember which
number has which index (the output requires knowing this), we pair
each k-gonal number n with its index k, i.e., we have (n,k) in the
list.

The above is what gonalsFrom does.

> type Gonal = Integer
> type Index = Int

> gonalsFrom :: Gonal -> Index -> [(Gonal,Index)]
> gonalsFrom s k = flip zip (repeat k) $ dropWhile (< s) $
>                  scanl1 (+) [1, fromIntegral k - 1 ..]

This merges two infinite lists in increasing order:

> merge :: Ord a => [a] -> [a] -> [a]
> merge xs@(x:xt) ys@(y:yt) = if x < y then x : merge xt ys
>                             else y : merge xs yt

Ha, the one-in-all function that solves the problem.  Given the
starting point and the list of indexes we are interested in:

1. compute the relevant lists of polygonal numbers
2. merge them into one single list in increasing order
3. group common polygonal numbers together
4. discard the rest
5. for each group, extract the common polygonal number and
   the indexes into a pair (number, [index1, index2, ...])
6. return five such pairs.  Done!

> poly :: Gonal -> [Index] -> [(Gonal,[Index])]
> poly s ks = take 5 $
>             map (\x -> (fst (head x), map snd x)) $
>             filter ((>= 2) . length) $
>             groupBy eq $
>             foldl1 merge $
>             map (gonalsFrom s) ks
>   where eq (x,_) (y,_) = x==y

> main = do n <- readLn
>           if n > 0
>              then do l <- getLine
>                      ks <- rreadIO (sequence (replicate n reader)) l
>                      s <- readLn
>                      putStr (unlines (map showAnswer (poly s ks)))
>                      main
>              else return ()
>   where showAnswer (n,ks) = show n ++ ':' :
>                             concat (intersperse " " (map show ks))