Решаем задачи Яндекс.Интервью в функциональном стиле

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

Зарегистрировавшись в системе, моё внимание сразу привлекла возможность решать задачи на Haskell. Дело в том, что я хоть и увлекаюсь программированием на этом языке, но не продвинулся дальше реализации задач из различных курсов образовательных on-line платформ. Решив, что их решение может оказаться интересным вызовом и повысит мой уровень, как разработчика, я приступил к их решению.

Кому интересно, что в итоге из этого вышло, добро пожаловать под кат.

A. Камни и украшения

Даны две строки строчных латинских символов: строка J и строка S. Символы, входящие в строку J, — «драгоценности», входящие в строку S — «камни». Нужно определить, какое количество символов из S одновременно являются «драгоценностями». Проще говоря, нужно проверить, какое количество символов из S входит в J.

Первая задача является разминочной, будем решать её «в лоб». Определим функцию jeweleryCount :: String -> String -> Int, которая с помощью свертки по переданному вторым аргументом списку просуммирует все случаи нахождения обрабатываемого элемента в первом списке. Для этих целей определим функцию elemInt на основе функции elem, которая в отличии от последней вернёт не True или False, а число 0 или 1. В функции main остаётся только считать две строки, передать их в соответствующую функцию и напечатать результат. Вердикт системы тестирования – OK, переходим ко второй задаче.

jeweleryCount :: String -> String -> Int
jeweleryCount j = foldr ((+).(elemInt j)) 0
                where elemInt s x = fromEnum $ elem x s

main :: IO ()
main = do
    j <- getLine
    s <- getLine
    print $ jeweleryCount j s

Исходный код решения этой и других задач также доступен в github-репозитории

B. Последовательно идущие единицы

Требуется найти в бинарном векторе самую длинную последовательность единиц и вывести её длину.

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

Для чтения входных данных определим функцию getUserInputs :: IO [Char], в которой сначала прочитаем число n – размер списка, а затем с помощью комбинатора replicateM получим функцию, которая n раз выполнит вызов функции head <$> getLine и объеденит полученные результаты в список.

import Control.Monad (replicateM)

onesCount :: [Char] -> Int
onesCount xs = onesCount' xs 0 0
    where
        onesCount' "" max curr 
            | max > curr = max 
            | otherwise  = curr
        onesCount' (x:xs) max curr
            | x == '1' = onesCount' xs max $ curr + 1 
            | curr > max = onesCount' xs curr 0 
            | otherwise = onesCount' xs max 0

getUserInputs :: IO [Char]
getUserInputs = do
    n <- read <$> getLine :: IO Int
    replicateM n $ head <$> getLine

main :: IO ()
main = do
    xs <- getUserInputs 
    print $ onesCount xs

Отправляем решение, вердикт — OK. Двигаемся дальше.

C. Удаление дубликатов

Дан упорядоченный по неубыванию массив целых 32-разрядных чисел. Требуется удалить из него все повторения.

Начнём с простой реализации. Определим функцию initial, которая считывает число, печатает его и возвращает завернутым в монаду IO. Также определим функцию deleteDoubles :: Int -> Int -> IO(), которая считывает число и печатает в его только в случае, если оно не равно второму аргументу (будем передавать туда число прочитанное на предыдущем шаге). После этого функция рекурсивно вызывает сама себя и таким образом переходит к следующему числу во входном потоке. Базой рекурсии является количество чисел, которое предстоит прочитать, будем передавать его первым аргументом.

import Control.Monad

initial :: IO Int
initial = do
    a <- read <$> getLine
    print a
    return a

deleteDoubles :: Int -> Int -> IO()
deleteDoubles 0 _ = return ()
deleteDoubles t a = do
    b <- read <$> getLine
    unless (a == b) $ print b 
    deleteDoubles (t-1) b

main :: IO ()
main = do
    t <- read <$> getLine
    unless (t < 1) $ initial >>= deleteDoubles (t-1)

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

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

import Control.Monad

deleteDoubles' _ [] = []
deleteDoubles' prev (x:xs)
    | prev /= x = x:(deleteDoubles' x xs)
    | otherwise = deleteDoubles' x xs

deleteDoubles (x:xs) = x:deleteDoubles' x xs

getUserInputs :: Int -> IO [Int]
getUserInputs t = replicateM t $ read <$> getLine

main :: IO ()
main = do
    t <- read <$> getLine
    unless (t < 1) $ (deleteDoubles <$> getUserInputs t) >>= mapM_ print

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

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

Функция, которая печатает или не печатает результат в зависимости от своих аргументов, после чего возвращает свой второй аргумент, завёрнутый в монаду IO, довольно проста, назовём её step:

step :: Int -> Int -> IO Int
step fst snd = unless (fst == snd) (print snd) >> return snd

С печатью или не печатью в зависимости от переданных значений мы разобрались, но как организовать чтение? Для этого воспользуемся функцией монадической свертки foldM :: (Foldable t, Monad m) => (b -> a -> m b) -> b -> t a -> m b, которую применим к списку функций чтения.
По типу функции foldM заметим, что на каждом шаге «распаковка» результата прошлого применения функции происходит под капотом самой foldM. Таким образом, на каждом шаге нам необходимо только запустить монадическое вычисление текущего элемента списка (по сути – прочитать следующее число) с помощью оператора bind (>>=) и вместе с предыдущим числом передать в step. В итоге получаем следующую программу

step :: Int -> Int -> IO Int
step fst snd = unless (fst == snd) (print snd) >> return snd

initial :: IO Int
initial = do
    a <- read <$> getLine
    print a
    return a

getUserInputs t = replicate t $ read <$> getLine

main :: IO ()
main = do
    t <- read <$> getLine
    unless (t < 1) $ do
        init <- initial
        foldM_ ((=<<) . step) init $ getUserInputs (t-1)

D. Генерация скобочных последовательностей

Дано целое число n. Требуется вывести все правильные скобочные последовательности длины 2 ⋅ n, упорядоченные лексикографически (см. https://ru.wikipedia.org/wiki/Лексикографический_порядок).
В задаче используются только круглые скобки.
Желательно получить решение, которое работает за время, пропорциональное общему количеству правильных скобочных последовательностей в ответе, и при этом использует объём памяти, пропорциональный n.

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

Определим рекурсивную функцию generate' :: Int -> Int -> [[Char]], которая вторым аргументом принимает количество скобок, которое ещё предстоит поставить, а первым – количество уже поставленных незакрытых открывающих скобок. Для шага рекурсии нам понадобится две вспомогательные функции: possible — возвращает список скобок, которые можно разместить на следующем шаге, и step – производит рекурсивный вызов функции generate' с необходимыми параметрами.

import Control.Monad(mapM_)

generate :: Int -> [String]
generate = generate' 0
        where
        generate' _ 0 = [[]]
        generate' a n = [x:xs | x <- possible, xs <- step x]
            where
            step '(' = generate' (a + 1) (n - 1) 
            step ')' = generate' (a - 1) (n - 1)
            possible
                | n == a = ")"
                | a == 0 = "("
                | otherwise = "()"

main :: IO ()
main = do
    n <- read <$> getLine
    let result = generate $ n * 2
    mapM_ putStrLn result

Отправляем решение, и понимаем, что мы не учли ограничение, которое накладывалось на используемое программой количество памяти – решение не проходит 14 тест из-за превышения лимита используемой памяти.

Модифицируем функцию generate' таким образом, чтобы она вместо конструирования всего списка правильных скобочных последовательностей сразу выводила их на экран. Для этого нам придётся добавить третий аргумент к функции – фрагмент последовательности, сконструированный к текущему шагу. Отмечу, что в данной реализации будем конструировать последовательность в обратном порядке – это позволит нам использовать конструктор списка (:) вместо более дорогостоящего оператора конкатенации (++).

import Control.Monad(mapM_)

generate :: Int -> IO()
generate = generate' "" 0
        where
        generate' xs _ 0 = putStrLn $ reverse xs
        generate' xs a n 
            | n == a = step ')'
            | a == 0 = step '('
            | otherwise = step '(' >> step ')'
            where
                step '(' = generate' ('(':xs) (a + 1) (n - 1) 
                step ')' = generate' (')':xs) (a - 1) (n - 1)

main :: IO ()
main = do
    n <- read <$> getLine
    generate $ n * 2

E. Анаграммы

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

Для решения этой задачи будем подсчитывать сколько раз в каждой строке встречается та или иная буква и сравнивать полученные результаты. Сразу понимаем, что стандартные списки нам не подходят, и необходимо использовать структуру данных, которая бы позволила эффективно обращаться к элементу по его индексу. Существует несколько типов данных, которые бы удовлетворяли нашим условиям, мы же воспользуемся стандартным неизменяемым массивом Data.Array (ещё существуют как минимум различные изменяемые массивы, а также Data.Vector).

Для конструирования необходимых массивов воспользуемся функцией hist :: (Ix a, Num b) => (a,a) -> [a] -> Array a b, которая по переданному списку элементов и диапазону, которому данные элементы должны принадлежать, формирует массив, который хранит в себе количество повторов элементов из списка. Данная функция хоть и не входит в модуль Data.Array, но часто приводится как пример использования другой, уже библиотечной функции accumArray. Нам остаётся только скопировать её реализацию и написать main – благо сравнение на равенство для Array Char Int уже определено. Обращаю ваше внимание на одну приятную особенность – в качестве индекса мы можем воспользоваться не только целыми числами, а любым представителем класса Ix. В нашем случае на эту роль естественным образом подходит Char.

import Data.Array

hist :: (Ix a, Num b) => (a,a) -> [a] -> Array a b
hist bnds is = accumArray (+) 0 bnds [(i, 1) | i<-is, inRange bnds i]

main = do 
    arr1 <- hist ('a','z') <$> getLine
    arr2 <- hist ('a','z') <$> getLine
    if (arr1 == arr2) then print 1 else print 0

F. Слияние k сортированных списков

Даны k отсортированных в порядке неубывания массивов неотрицательных целых чисел, каждое из которых не превосходит 100. Требуется построить результат их слияния: отсортированный в порядке неубывания массив, содержащий все элементы исходных k массивов.
Длина каждого массива не превосходит 10 ⋅ k.
Постарайтесь, чтобы решение работало за время k ⋅ log(k) ⋅ n, если считать, что входные массивы имеют длину n.

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

merge :: [Int] -> [Int] -> [Int]
merge [] ys = ys
merge xs [] = xs
merge (x:xs) (y:ys)
    | x < y     = x:merge xs (y:ys)
    | otherwise = y:merge (x:xs) ys

Хорошо, мы умеем производить слияние двух списков. А что нам делать со списком списков? Выполнить его свёртку с этой функцией! Таким образом мы объединим все списки в один, и нам останется его только распечатать.

Решение

import Control.Monad

merge :: [Int] -> [Int] -> [Int]
merge [] ys = ys
merge xs [] = xs
merge (x:xs) (y:ys)
    | x < y     = x:merge xs (y:ys)
    | otherwise = y:merge (x:xs) ys

mergeLists :: [[Int]] -> [Int]
mergeLists = foldl merge []

getUserInputs :: Int -> IO [[Int]]
getUserInputs t = replicateM t $ do
    n <- getLine
    return $ tail $ read <$> words n

main :: IO ()
main = do
    k <- read <$> getLine
    lists <- getUserInputs k
    let res = mergeLists lists
    mapM_ (putStrLn . show) res

Однако, у этого решения имеются две серьёзные проблемы – вычислительная сложность оказывается выше требуемой — O(k^2 ⋅ n) вместо О(k ⋅ log(k) ⋅ n), плюс ко всему оно использует довольно много дополнительной памяти. Как итог, такое решение проваливает тест номер 17 из-за превышения лимита используемой памяти — 17.27Mb вместо разрешенных 10Mb.

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

Следующим шагом попробуем реализовать подход, который был предложен в исходной статье с разбором данных задач. Напомню, он основывается на использовании структуры данных, предоставляющей эффективный способ извлечения минимального элемента. В качестве такой структуры выберем Data.Set. Инициализируем Set списком первых элементов, затем на каждом шаге будем извлекать и печатать минимальный элемент, после чего добавлять следующий элемент из соответствующего списка. Кроме этого, нам понадобится структура Data.Sequence для хранения самих списков. Она была выбрана из соображений, что на каждом шаге необходимо как иметь быстрый доступ к списку по его индексу (что не может обеспечить список), так и изменять элемент этот элемент без необходимости копирования всей структуры (что в общем случае не может обеспечить неизменяемый Data.Array).

Таким образом, имеем следующую программу:

Решение

import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import Control.Monad
import Data.Foldable

mergeLists :: Set.Set (Int, Int) -> Seq.Seq [Int] -> IO ()
mergeLists set seq 
    | Set.null set = return ()
    | otherwise  = do
        let ((val, idx), set') = Set.deleteFindMin set
        print val
        if null (Seq.index seq idx)
            then mergeLists set' seq
            else mergeLists (Set.insert (head (Seq.index seq idx), idx) set') (Seq.adjust tail idx seq)

getUserInputs :: Int -> IO [[Int]]
getUserInputs t = replicateM t $ do
    n <- getLine
    return $ tail $ read <$> words n

main :: IO ()
main = do
    t <- read <$> getLine
    lists <- getUserInputs t
    let init_seq = Seq.fromList (filter (not . null) lists)
    let init_heap = Set.fromList (zipWith (,) (toList (fmap head init_seq)) [0..])
    mergeLists init_heap $ tail <$> init_seq

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

Мы уже выполняли подсчёт количества входящих символов при решении предыдущей задачи об анаграммах. Также, как и при её решении, воспользуемся Data.Array. Для начала реализуем функцию addToArray :: Array Int Int -> [Int] -> Array Int Int, которая формирует массив на основе существующего путём увеличения значений по индексам, которые соответствуют значениям из списка.

addToArray :: Array Int Int -> [Int] -> Array Int Int
addToArray acc elems = accum (+) acc [(i, 1) | i<-elems]

Затем, воспользуемся подходом, известным нам по задаче об удалении повторов — с помощью монадической свёртки последовательно применим функцию addToArray к k исходным массивам. И… получаем всё тот же результат 10.26Mb на 17 тесте. И тут самое время вспомнить что foldl (аналогом которого является foldM) согласно принятому порядку редукции сначала развернёт всю цепочку вложенных выражений и только потом приступит к их активному вычислению. Как известно, для борьбы с этим фактом в модуле Data.List реализована функция foldl', использующая функцию seq :: a -> b -> b, которая сначала приводит первый аргумент в слабую головную нормальную форму, то есть редуцирует до получения внешней части — значения функции или конструктора, а затем возвращает второй (https://www.ibm.com/developerworks/ru/library/l-haskell4/index.html). Нам же ничего не остается делать кроме того как самостоятельно реализовать функцию foldM'.

foldM' :: (Monad m) => (a -> b -> m a) -> a -> [b] -> m a
foldM' _ z [] = return z
foldM' f z (x:xs) = do
    z' <- f z x
    z' `seq` foldM' f z' xs

В результате количество используемой памяти на 17 тесте снизилось почти в два раза и составило 5.64Mb! И хотя 17 и 18 тесты были успешно пройдены, данная реализация не проходила уже 19 тест по той же причине превышения лимита использования памяти — 10.25Mb.

Окей, идём дальше — мы ещё не попробовали Data.Array.Unboxed. Этот вид массивов примечателен тем, что, в отличии от стандартного, своими элементами может хранить сами значения, а не указатели на них (https://wiki.haskell.org/Arrays#Unboxed_arrays). Благодаря этому, такие массивы занимают меньшее пространство в памяти и более производительны. Для того чтобы ими воспользоваться нам необходимо только поменять импорт и типы функций, так как Data.Array и Data.Array.Unboxed реализуют один интерфейс неизменяемых массивов IArray.

Отправляем решение – потребление памяти снизилось в 4.5 раза до 2,26 MB, но оно не прошло уже ограничение по времени — время исполнения составило 1.09 секунды. С чем это может быть связано? Судя по тому, что время исполнения остальных тестов осталось прежним, думаю, что причина не в том, что unboxed-массив оказался медленнее boxed, а в особенности системы тестирования. Похоже, выполнение задачи прерывается, как только нарушено одно из ограничений. Однако, в очень редких случаях эта реализация всё-таки проходит 19 тест с результатом 0.98 секунды, но заваливает тест номер 20 также из-за превышения лимита времени.

После этого я попробовал воспользоваться unsafe аналогом функции accum, которая в теории должна быть быстрее, различными способами буферизации (функция hSetBuffering :: Handle -> BufferMode -> IO ()) и изменяемыми массивами IOArray, но не один из этих способов не принёс никаких результатов.

Я не склонен считать, что лимиты для Haskell заданы слишком жестко, и надеюсь, что всё-таки существует решение, которое пройдёт все тесты. В репозиторий проекта я выложил несколько различных версий кода решения этой задачи (с Array и IOArray), возможно это станет отправной точкой для решения, которое пройдёт уже все тесты.

Заключение

Даже несмотря на то, что по итогу мне поддались только пять задач из шести, свою главную задачу – попрактиковаться в функциональном программировании, я выполнил. Не последнюю роль в этом сыграли и жесткие ограничения на потребляемые программой ресурсы, которые вынуждали искать всё новые и новые подходы к решению задач. Надеюсь их описание будет полезно для тех, кто только начинает свой путь в функциональном программировании

Оказался ли функциональный подход удобен для решения подобного рода задач? Честно говоря, у меня осталось двоякое впечатление. С одной стороны, решения большинства задач оказывались весьма лаконичными, и не последнюю роль в этом сыграли выразительные средства самого Haskell, а также его богатая стандартная библиотека. С другой, нельзя не признать, что в большинстве случаев управление потребляемыми ресурсами может представлять собой определенную проблему, что не позволит решить задачу в заданных ограничения.