본문 바로가기

Haskell

원소 필터링 문제

연습문제 - 이번에는 약간 쉽지 않다. N개의 정수로 이루어진 리스트 A를 받는다. 그리고 정수 K를 받는다. 리스트 A에서 K번 이상 등장한 정수들을 필터링하여, 최초에 등장했던 순서대로 출력하는 것이다.

예를 들어 A = [4, 5, 2, 5, 4, 3, 1, 3, 4] 이고, K=2라면 이 중에서 두 번 이상 등장한 원소는 순서대로 4, 5, 3 이므로 출력은 4 5 3 이다. 만약 K번 이상 등장한 원소가 없다면 -1을 출력해야 한다.

출처 : https://www.hackerrank.com/challenges/filter-elements/problem

입력

  • 첫줄에서는 시행횟수를 받는다
  • 각 시행의 첫입력은 리스트의 길이와 K 값을 공백으로 구분하여 받음
  • 각 시행의 둘째 입력은 리스트의 내용

접근

원소의 등장 순서를 유지하면서 등장 횟수를 같이 셀 수 있는 CountedList 타입을 만들어서 이 문제를 풀어보자. CountedList 타입은 리스트와 유사한 재귀적인 구조로 되어 있는데, 각 노드의 값과 카운트를 동시에 가지고 있다.

우선 임의의 타입 a를 원소로 가지는 CountedList를 다음과 같이 정의해보자.


data CountedList a =
	Node a Int (CountedList a) | --재귀적으로 중첩
	Nil  -- 끝을 나타냄
	deriving (Show) -- 출력을 위해.

리스트에 원소를 추가하는 append 함수를 작성해보자. 원소의 타입인 a는 같은지 비교가 가능해야 하기 때문에 Eq의 인스턴스여야 한다.


append :: (Eq a) => CountedList a -> a -> CountedList a
append Nil nval = Node nval 1 Nil -- 빈리스트에 추가하는 경우
append (Node val cnt next) nval
	-- 원소값이 같으면 카운트만 1올림
	| val == nval = Node val (cnt+1) next 
	-- 원소값이 다르면 뒤쪽에 추가
	| otherwise = Node val cnt (append next nval)

자 그러면 리스트 A에 대해서 빈 카운트리스트에 차곡차곡 밀어넣으면 변환이 완성된다.


lst2clst :: (Eq a) => [a] -> CountedList a
lst2clst = foldl append Nil

그렇다면 역으로 카운트리스트에서 K회 이상 나온 원소만 골라내보면 다음과 같다.


filtered :: Int -> CountedList a -> [a]
filtered _ Nil = []
filtered k (Node val cnt next) = 
	if k < cnt then filtered k next
	           else val:(filtered k next)

문제 풀이를 위한 구성함수가 모두 준비되었으니, 자 인제 문제를 풀어보자.


-- forM_을 써야하니
import Control.Monad

main :: IO ()
main = 
	let readInts = (map (read::String->Int) 
					. words) <$> getLine
    n <- readLn :: IO Int
    forM_ [1..n] $ \_ -> do
      (_:k:_) <- readInts
      (lst) <- readInts
      let result = filtered k . lst2clst $ lst
      if null result 
      	then do putStrLn "-1"
      	else do putStrLn . unwords . map show $ result
      

최종소스는 https://gist.github.com/stripe-q/c436596982d9e9402027b0963f9a2043 에 있다.

개선

문제는 이게 간단한 리스트에 대해서는 잘 작동하는데, 큰 데이터인 경우에는 처참한 수준으로 느리다는게 문제이다. 첨부된 input04.txt를 대상으로 실행해본 결과 거의 15초 수준으로 수행시간이 걸린다. (물론 데이터가 무식한 수준으로 많기는 하다)


$ time ./h016 < input04.txt
real    0m13.959s
user    0m0.015s
sys     0m0.015s

이 문제를 재귀적으로 풀어서는 답이 없을 것 같고 "특단의 조치"를 취해야 할 것 같다.

리스트를 정렬한 후에 map head . filter (>= k).length . group을 적용하면 간단하게 k 번 이상 나온 원소들을 추릴 수 있다. 하지만 이렇게 정렬을 하는 경우 "원래의 등장순서"라는 정보를 잃게된다.

그렇다면 "원래의 등장순서"를 유지해주게 하면 되지 않을까? zip ns [1..]을 적용해주면 (원소, 인덱스)의 쌍이 만들어진다.

  1. (원소, 인덱스)의 쌍을 만든다.
  2. 이를 원소 순으로 정렬한 후
  3. 이를 원소 기준으로 그룹화한다.
  4. 조건에 맞게 필터링한 후
  5. 다시 인덱스순으로 정렬한다.

이 과정을 거치면 좀 더 빠르게 풀 수 있을 것 같다. (보통 리스트의 정렬은 NlogN 이고 재귀의 경우 최악의 상황에서 N^2 가량 되므로 이쪽이 훨씬 이득)

예를들어서 zip ns [1..] 을 적용하면 [(4,1), (1,2),(2,3),...]과 같은식의 리스트로 변환된다. 즉 각 원소는 처음의 위치정보와 튜플로 묶이게 되었다. 이를 다시 값 기준으로 정렬하기 위해서는 sortBy 함수를 쓴다. 이 함수의 타입 시그니처는 다음과 같다.


sortBy :: (a -> a -> Ordering) -> [a] -> [a]

비교하고자하는 값은 (Int, Int) 타입이고 이 때 튜플의 첫번째 요소를 비교해야 한다. 따라서 정렬코드는 다음과 같다.


sorted = sortBy (\x y -> (fst x) `compare` (fst y)) ...

그런데 이 정렬용 비교함수가 어째 마뜩찮아 보인다. (좀 귀찮기도 하고) 이 때 쓰라고 있는 함수가 on (Data.Function에 있다.) 이다.

이 함수의 시그니처는 다음과 같다.


on :: (b -> b -> c) -> (a -> b) -> a -> a -> c

(*) 'on' f\x y -> f x * f y가 된다. 즉 왼쪽에 오는 2인자 함수의 각 인자에 대해서 오른쪽에 오는 함수를 먼저 적용하는 일종의 분배에 의한 합성 함수가 된다. 따라서 앞서 sortBy에 사용했던 비교 함수는 다음과 같이 변형할 수 있다.


sorted = sortBy (compare `on` fst) index

OK? 이제 빌딩 블럭들이 모두 모였으니 새로운 코드를 짜 보도록 하자.


import Data.List -- groupBy, sortBy
import Data.Function -- on
import Control.Monad -- forM


process :: Int -> [Int] -> [Int]
process k ns =
  let indexed = zip ns [1..]
      sorted = sortBy (compare `on` fst) indexed
      grouped = groupBy ((==) `on` fst) sorted
      filtered = filter ((>= k) . length) grouped
  in result = map fst . sortBy (compare `on` snd) .
  			  map head $ filtered
  			  
main = 
	let readInts = (map (read::String->Int) 
					. words) <$> getLine
    n <- readLn :: IO Int
    forM_ [1..n] $ \_ -> do
      (_:k:_) <- readInts
      (lst) <- readInts
      let result = process k lst
      if null result 
      	then do putStrLn "-1"
      	else do putStrLn . unwords . map show $ result
      

컴파일하고 시간을 재보면 큰 향상이 있다! 1초대에 해결됐다.


$ time ./h016-2 < input04.txt
...
real    0m1.087s
user    0m0.031s
sys     0m0.000s

의문

이 문제가 있는 카테고리는 "Recursion"인데, 왜 재귀로 풀어서는 답이 없을까? 만약 파이썬이라면 사전과 같은 맵 구조를 이용해서 풀었을 것이다. 하스켈에서도 맵 구조가 있다고 하는데 (아니, immutable 한 언어에서 이게 가능함?) 어떤 식으로 구현되는것인지 한 번 공부해봐야 할 듯 하다.