Clojure Tail Recursion with Prime Factors

Question

I'm trying to teach myself clojure and I'm using the principles of Prime Factors Kata and TDD to do so.

Via a series of Midje tests like this:

(fact (primefactors 1) => (list))

(fact (primefactors 2) => (list 2))

(fact (primefactors 3) => (list 3))

(fact (primefactors 4) => (list 2 2))

I was able to create the following function:

(defn primefactors 
    ([n] (primefactors n 2))
    ([n candidate] 
        (cond (<= n 1) (list)
              (= 0 (rem n candidate)) (conj (primefactors (/ n candidate)) candidate)
              :else (primefactors n (inc candidate))
        )
    )
)

This works great until I throw the following edge case test at it:

(fact (primefactors 1000001) => (list 101 9901))

I end up with a stack overflow error. I know I need to turn this into a proper recur loops but all the examples I see seem to be too simplistic and only point to a counter or numerical variable as the focus. How do I make this recursive?

Thanks!

Answer 1

Here's a tail recursive implementation of the primefactors procedure, it should work without throwing a stack overflow error:

(defn primefactors 
  ([n] 
    (primefactors n 2 '()))
  ([n candidate acc]
    (cond (<= n 1) (reverse acc)
          (zero? (rem n candidate)) (recur (/ n candidate) candidate (cons candidate acc))
          :else (recur n (inc candidate) acc))))

The trick is using an accumulator parameter for storing the result. Notice that the reverse call at the end of the recursion is optional, as long as you don't care if the factors get listed in the reverse order they were found.

Answer 2

Your second recursive call already is in the tail positions, you can just replace it with recur.

(primefactors n (inc candidate))

becomes

(recur n (inc candidate))

Any function overload opens an implicit loop block, so you don't need to insert that manually. This should already improve the stack situation somewhat, as this branch will be more commonly taken.

The first recursion

(primefactors (/ n candidate))

isn't in the tail position as its result is passed to conj. To put it in the tail position, you'll need to collect the prime factors in an additional accumulator argument onto which you conj the result from the current recursion level and then pass to recur on each invocation. You'll need to adjust your termination condition to return that accumulator.

Answer 3

The typical way is to include an accumulator as one of the function arguments. Add a 3-arity version to your function definition:

(defn primefactors
  ([n] (primefactors n 2 '()))
  ([n candidate acc]
    ...)

Then modify the (conj ...) form to call (recur ...) and pass (conj acc candidate) as the third argument. Make sure you pass in three arguments to recur, i.e. (recur (/ n candidate) 2 (conj acc candidate)), so that you're calling the 3-arity version of primefactors.

And the (<= n 1) case need to return acc rather than an empty list.

I can go into more detail if you can't figure the solution out for yourself, but I thought I should give you a chance to try to work it out first.

Answer 4

This function really shouldn't be tail-recursive: it should build a lazy sequence instead. After all, wouldn't it be nice to know that 4611686018427387902 is non-prime (it's divisible by two), without having to crunch the numbers and find that its other prime factor is 2305843009213693951?

(defn prime-factors
  ([n] (prime-factors n 2))
  ([n candidate]
     (cond (<= n 1) ()
           (zero? (rem n candidate)) (cons candidate (lazy-seq (prime-factors (/ n candidate)
                                                                              candidate)))
           :else (recur n (inc candidate)))))

The above is a fairly unimaginative translation of the algorithm you posted; of course better algorithms exist, but this gets you correctness and laziness, and fixes the stack overflow.

Answer 5

A tail recursive, accumulator-free, lazy-sequence solution:

(defn prime-factors [n]
  (letfn [(step [n div]
            (when (< 1 n)
              (let [q (quot n div)]
                (cond
                  (< q div)           (cons n nil)
                  (zero? (rem n div)) (cons div (lazy-step q div))
                  :else               (recur n (inc div))))))
          (lazy-step [n div]
            (lazy-seq
              (step n div)))]
    (lazy-step n 2)))

Recursive calls embedded in lazy-seq are not evaluated before iteration upon the sequence, eliminating the risks of stack-overflow without resorting to an accumulator.