Решение задачи J финальной части конкурса ACM ICPC 2011

Всем привет. Давно не брал я в руки шашек (в смысле, не писал ничего полезно-тематического). Сегодня мы с вами попробуем решить одну из задач, представленных в финальной части чемпионата мира по программированию среди студентов ACM ICPC 2011. С условиями задач вы можете ознакомиться здесь http://cm.baylor.edu/digital/icpc2011.pdf

Решать будем задачу J (problem J), потому как она мне сразу понравилась своей простотой )) В качестве языка программирования будем использовать Haskell (вообще, насколько я понял, на конкурсе принимались решения только на C++ или Java, но так как мы с вами в конкурсе не участвуем, позволим себе такую вольность).

Итак, в чем суть задачи вкратце. У нас есть N кирпичей-блоков, из которых можно строить пирамиды, где каждый уровень представляет собой объем из KxKx1 кирпичей-блоков, проекция которого на основание пирамиды есть квадрат со стороной K.
Пирамиды могут быть 2ух типов: высокие, когда на каждом следующем после основания уровне K уменьшается на 1 (например, в основании лежит 5x5 кирпичей, затем на втором уровне 4x4 и т.д., вверху будет всегда 1 кирпич), и низкие, когда на каждом следующем уровне K уменьшается не на 1, а на 2 (например, в основании 5x5 кирпичей, затем на втором уровне - 3x3 кирпичей, а на третьем 1). Или такой вариант низкой пирамиды: 6x6,4x4,2x2. Понятно, что низкая пирамида на самом последнем (верхнем) уровне может иметь либо 1, либо 2x2 кирпичей.

Так вот наша задача: имея N кирпичей, построить из них такое количество пирамид, чтобы, во-первых, израсходовать все кирпичи (если это невозможно, нужно возвращать специальный вариант ответа), а во-вторых, чтобы количество получаемых при этом пирамид было минимальным. Кроме того, есть дополнительные условия: все строящиеся пирамиды должны быть различными, каждая пирамида должна иметь минимум 2 уровня высоты, в полученном списке пирамид каждая пирамида должны состоять из максимально возможного количества кирпичей, затраченных на её постройку.

Входные данные представляют собой список чисел по одному в каждой строке - количество кирпичей для заданного случая. В последней строке содержится 0 - признак окончания обработки. На выход для каждого случая нужно выдавать решения. Например, вот для таких входных данных
====
29
28
0
====
наша программа должна выдать в ответ:
===
Case 1: 3H 3L 2H
Case 2: impossible
==
, то есть, когда у нас есть 29 кирпичей, решением будет строительство 3-х пирамид: одной высокой (High) со стороной основания 3 (3x3,2x2,1x1), одной низкой со стороной основания 3 (3x3,1x1) и одной высокой со стороной основания 2 (2x2,1x1). Сумма всех кирпичей, если вы проверите, составит ровно 29. Для варианта же 28 кирпичей решения не существует. Да, еще один момент: каждое из входных данных (количество кирпичей) может находиться в диапазоне от 1 до 10^6 включительно.

Итак, приступим.
Определим тип пирамиды как "H","L1" и "L2". Почему я разделил "низкие" пирамиды на 2 подтипа ? Просто так нам будет легче с ними работать. Пирамида типа L1 будет иметь четную ширину основания (то есть, на самом верхнем уровне будет 2x2 кирпичей), а пирамида L2 - нечетную (таким образом, на самом верхнем уровне у неё будет 1 кирпич, как и у "высокой" пирамиды).

import Data.List (sortBy)

data PType = L1 | L2 | H deriving (Eq)

Пирамиду мы представляем кортежем-тройкой, где первым элементом идет высота пирамиды, вторым - её тип (H|L1|L2), а третьим - для удобства - будем хранить количество кирпичей, которое требуется для её постройки.

-- (высота, тип пирамиды, количество кирпичей)
type Pyr = (Int,PType,Int) 

getBricks :: Pyr -> Int
getBricks (_,_,br) = br

Тогда для каждого типа пирамиды зададим бесконечный список, представляющий количество кирпичей на каждом уровне, начиная с самого верха.

getLst :: PType -> [Int]
getLst H  = [1,2..]
getLst L1 = [2,4..]
getLst L2 = [1,3..]

После этого определим функцию, которая генерирует для заданного количества кирпичей nBr (входное значение) список всех возможных пирамид, высота которых не меньше 2-ух и количество кирпичей в которых меньше либо равно заданному nBr.

genPyrLst :: Int -> [Pyr]
genPyrLst nBr = 
    let gen res acc t lvl lst = 
        let hSq = (head lst) ^ 2 in 
            if acc + hSq > nBr then res
            else
               if lvl > 0 then gen ((lvl+1,t,acc+hSq):res) 
                                   (acc + hSq) 
                                    t 
                                   (lvl + 1) 
                                   (tail lst) 
               else gen res (acc + hSq) t 
                            (lvl + 1) (tail lst) 
 in (gen [] 0 H  0 (getLst H)) ++ 
    (gen [] 0 L1 0 (getLst L1)) ++
    (gen [] 0 L2 0 (getLst L2))

Итак, на данном этапе у нас есть список всех возможных пирамид, которые можно построить из заданного количества кирпичей nBr.

Теперь мы перешли к, на мой взгляд, ключевому моменту алгоритма решения задачи.
Нам необходимо перебрать все возможные варианты сочетаний этих пирамид (помним, что по условию задачи все пирамиды должны быть разными), при этом соблюдая 2 условия:
1. если возможно несколько решений, то выбираем то, в котором количество пирамид меньше (A)
2. если возможно несколько решений, то выбираем то, для которого попирамидное сравнение решений по количеству кирпичей в пирамиде даст "больше" (B).

Как это можно реализовать? Ну, для начала мы упорядочим наш список пирамид по убыванию количества кирпичей, требуемых для их постройки. Затем будем формировать списки всех возможных сочетаний пирамид из нашего общего списка длины N сначала по 1, затем по 2, затем по 3 и т.д. И тут же на каждом этапе будем проверять, а не найдено ли уже решение нашей задачи, и если найдено, то сразу возвращать результат (вернуть надо лишь один - правильный - вариант). Таким образом, мы решаем первую проблему (см.A), так как если существует несколько вариантов решений, то первым нам попадет то, которые содержит меньшее количество пирамид в списке. Для борьбы со второй проблемой (см.B) будем формировать решения таким образом, чтобы каждый следующий элемент нашего списка содержал в себе пирамиды, количество кирпичей которых _не больше_ количества кирпичей, содержащихся в предыдущих элементах списка (всё это, конечно, в рамках конкретной глубины рекурсии). Как мы этого добьемся ?
Как я уже написал выше, отсортируем наш список пирамид в порядке убывания количества кирпичей, требуемых для их постройки.

Рассмотрим на примере списка чисел. Предположим, что у нас есть список чисел [5,3,2,1]. Список всех возможных подсписков длиной 1 представляет собой список списков [[5],[3],[2],[1]]. Теперь сформируем по этому полученному списку списков список списков длины 2. Что для этого надо сделать ? Пройдем по нашему списку списков и для каждого элемента составим новый список списков, состоящий из этого элемента и всех _последующих за этим элементом_ элементов списка. То есть, берем [5]. Следующие элементы за ним - [3],[2],[1]. Тогда подсписки будут [[5,3],[5,2],[5,1]]. Для элемента [3] получаем список [[3,2],[3,1]] и для [2] -> [[2,1]].
Склеив эти списки, получим [[5,3],[5,2],[5,1],[3,2],[3,1],[2,1]].

Легко заметить, что чем ближе элементы к началу списка, тем бОльшие числа они содержат в самом начале. И проверяя эти списки слева направо, мы всегда первым найдем правильное решение для заданной длины подсписков (в данном случае, для длины 2). Понятно, что самый худший вариант - когда решения нет. В этом случае нам придется перебрать все возможные варианты всех допустимых длин.

Применительно к нашей задаче для ускорения процесса вычисления (правда, за счет увеличения роста используемой памяти) будем в качестве элементов списка результатов хранить не просто список пирамид, а пару-кортеж, где первым элементом пары будет список пирамид результата, а вторым элементом - тот самый список оставшихся пирамид для генерации элементов следующего уровня рекурсии.

Кроме того, чтобы исключить обработку ненужных вариантов сочетаний (тех, в которых общая сумма кирпичей на конкретном этапе превысила исходное количество кирпичей, доступных для строительства всех пирамид), мы будем фильтровать списки после каждого этапа обработки.

Получаемый код выглядит так:

isFind :: [([Pyr],[Pyr])] -> Int -> [([Pyr],[Pyr])]
isFind [] n = []
isFind ((lst,lTail):xs) n = 
  let summ l = sum (map getBricks l) in
  if summ lst == n then [(lst,lTail)] else isFind xs n 

procOne ::   [Pyr] -> [Pyr] -> [([Pyr],[Pyr])]
procOne accElem l = f l []
  where
     f l res = 
          if l == [] then reverse res 
          else f (tail l) 
                 ((accElem ++ [head l],tail l):res)

process :: [([Pyr],[Pyr])] -> Int -> Int -> [([Pyr],[Pyr])]
process l 0 _  = []
process [] _ _ = []
process l n nBr = 
 let doStep acc = 
            concatMap (\(a,l) -> procOne a l) acc in
 let isCorrect n (lst,lTail) = 
            sum (map getBricks lst) <= n in
 let res = filter (isCorrect nBr) (doStep l) in 
 let srch = isFind res nBr in 
 if srch == [] then process res (n-1) nBr
 else srch

Далее, нам нужно написать функцию, которая будет выводить наш результат в том виде, который требуется в условии. Требуемая функция drawPyrLst вкупе с сопутствующими функциями выглядит так:

drawType :: PType -> String
drawType H = "H" 
drawType _ = "L" 

drawPyramid :: Pyr -> String
drawPyramid (h,t,br) = let base = (getLst t) !! (h-1) in
                       show base ++ drawType t

drawPyrLst :: [Pyr] -> String
drawPyrLst l = foldl (\a el -> a ++ 
                     (if a == [] then "" else " ") ++ 
                      drawPyramid el) [] l

Ну и на заключительном этапе оформим один шаг эксперимента отдельной функцией:

processStep :: Int -> Int -> IO ()
processStep nBr x = 
    let sortByBricks p1 p2 = 
        if getBricks p1 > getBricks p2 then LT else GT in
    let l = sortBy sortByBricks (genPyrLst nBr) in 
    let lstP = process [([],l)] (length l) nBr in
        if lstP == [] then 
             putStrLn $ "Case " ++ show x ++ ": impossible"
        else let (l2,tl)  = head lstP in
             putStrLn $ "Case " ++ 
                         show x ++ ": " ++ drawPyrLst l2

На вход она получает исходное количество кирпичей и номер текущего эксперимента. функция выполняет расчет для заданных входных данных и выводит в stdout результат: либо строку о том, что решения не существует, либо собственно, найденный список пирамид в требуемом формате.

Ну и сама функция main, которая считывает данные из stdin и вызывает требуемое количество раз функцию processStep с требуемыми параметрами:

main :: IO ()
main = 
        let processAll cntr = do
                x <- getLine
                let xInt = read x::Int
                if xInt == 0 then return () 
                else do
                        processStep xInt cntr
                        processAll (cntr+1)
        in processAll 1

Вот, собственно, и всё. Для примера, вот на таких входных данных:
===
29
28
48
51
20
40
50
60
1000
30000
780000
992000
0
===
exe-файл, собранный в Haskell Platform 2011.2.0.1 под Windows, выдает вот такие результаты
==
Case 1: 3H 3L 2H
Case 2: impossible
Case 3: impossible
Case 4: impossible
Case 5: 4L
Case 6: 5L 2H
Case 7: 4H 4L
Case 8: 5H 2H
Case 9: 15L 9H 5L
Case 10: 55L 13L 9H
Case 11: 132H 23H 11L
Case 12: 180L 24L 18L
==
за время порядка 100-150 миллисекунд.

Надеюсь, что данная заметка будет кому-нибудь интересна.