n-th smallest number in a matrix where the row and column are monotonically increasing

Smallest.hs

{-# LANGUAGE ScopedTypeVariables #-}
module Smallest where 

type List a = [ a]

merge :: Ord a => List a -> List a -> List a
merge as bs = 
    case (as, bs) of 
        ([], _) -> bs
        (_, []) -> as
        (x:xs, y:ys) -> 
            if x == y then 
                x : merge xs ys

            else if x < y then 
                x : merge xs (y:ys)

            else
                y : merge (x:xs) ys


{-| a helper to prevent an infinite loop. For the `foldr` below to be productive (i.e. actually produce values instead
of recursing endlessly) we have to let it produce something. Based on the assumption that 
the row and the column are increasing, the first value from the first argument (`xs`) is always the smallest. 
-}
mergeHelper :: Ord a => List a -> List a -> List a
mergeHelper [] ys = ys
mergeHelper (x:xs) ys = 
        x : merge xs ys

-- infinite row 1,2,3,4,5
row :: List Int
row = [1..]

-- infinite matrix
--
-- 1 2 3 4
-- 2 3 6 8
-- 3 6 8 12
matrix :: List (List Int)
matrix = 
    map (\n -> map (*n) [1..]) row

{-| recursively merge (i.e. take the first two rows, merge them, take the next row, merge it with the result
-}
sorted :: List (List Int) -> List Int
sorted = foldr mergeHelper [] 


main :: IO ()
main = do
    -- parse an Int from standard input
    n :: Int <- read <$> getLine
    
    -- get the n-th item from the sorted matrix
    let answer = (sorted matrix) !! n
    print answer

smallest.py

import itertools


def merge(xs, ys):
    while True:
        # try to pop one element from xs, 
        # yield all the elements of ys if that fails
        try:
            # tee copies an iterator
            (origional_xs, xs) = itertools.tee(xs)
            x = next(xs)
        except StopIteration:
            yield from ys
            return

        # try to pop one element from ys, 
        # yield all the elements of xs if that fails
        try:
            (origional_ys, ys) = itertools.tee(ys)
            y = next(ys)
        except StopIteration:
            yield from xs
            return

        if x == y:
            yield x

        elif x < y:
            yield x
            ys = origional_ys

        else:
            yield y
            xs = origional_xs


def force(thunk):
    return thunk()


def mergeHelper(xs, ys):
    try:
        x = next(xs)
    except StopIteration:
        yield from force(ys)
    else:
        yield x
        yield from merge(xs, force(ys))


def generateMatrix():
    start = 2
    for r in itertools.count(start):
        yield (r * n for n in itertools.count(start))


matrix = generateMatrix()


def foldr(function, default, list):
    try:
        x = next(list)
    except StopIteration:
        yield from default
    else:
        # make sure the recursive call to foldr is not evaluated (thus lazy)
        # by wrapping it into a function
        # mergeHelper will force evaluation when needed
        yield from function(x, lambda: foldr(function, default, list))


def result(items):
    yield from foldr(mergeHelper, iter([]), items)


if __name__ == '__main__':
    v = result(matrix)
    v2 = itertools.islice(v, 20)
    print(list(v2))