(module 3n-plus-1-problem mzscheme

  ;;  A solution to the 3n+1 (Collatz numbers) problem.
  ;;  http://acm.uva.es/p/v1/100.html
  ;;
  ;;  Danny Yoo (dyoo@cs.wpi.edu / dyoo@hkn.eecs.berkeley.edu)
  ;;
  
  (require (only (lib "etc.ss") local identity))
  (provide collatz find-max-chain)
  
  
  ;; We take advantage of memoization to trade off memory for speed.
  ;;
  ;; memoize: (natural -> natural) -> (natural -> natural)
  ;; Produces a memoized version of f that remembers previous
  ;; computations.
  (define (memoize/vector f)
    (local ((define *MAXIMUM-CACHED* 1000000)
            (define table (make-vector *MAXIMUM-CACHED* #f))
            (define (in-domain? n)
              (< n *MAXIMUM-CACHED*))
            (define (lookup n)
              (if (in-domain? n)
                  (vector-ref table n)
                  #f))
            (define (store! n val)
              (when (in-domain? n)
                (vector-set! table n val))))
      (lambda (var)
        (cond [(lookup var) => identity]
              [else
               (let ([result (f var)])
                 (store! var result)
                 result)]))))
  
  
  ;; memoize/hash: (natural -> natural) -> (natural -> natural)
  ;; Another data structure we can use is a hash-table.
  (define (memoize/hash f)
    (local ((define *MAXIMUM-CACHED* 1000000)
            (define table (make-hash-table 'equal))
            (define (in-domain? n)
              (< n *MAXIMUM-CACHED*))
            (define (lookup n)
              (hash-table-get table n #f))
            (define (store! n val)
              (when (in-domain? n)
                (hash-table-put! table n val))))
      (lambda (var)
        (cond [(lookup var) => identity]
              [else
               (let ([result (f var)])
                 (store! var result)
                 result)]))))
  
  
  ;; We provide a convenient syntax for defining memoizing functions
  ;; that should look similiar to regular DEFINE.
  (define-syntax (define/memo stx)
    (syntax-case stx ()
      [(_ (name var) memoizer body ...)
       (syntax/loc stx
         (define name (memoizer (lambda (var) body ...))))]))
  
  
  ;; collatz: number -> number.
  (define/memo (collatz n) memoize/vector
    (cond [(= n 1) 1]
          [(even? n)
           (add1 (collatz (/ n 2)))]
          [else
           (+ 2 (collatz (/ (+ 1 (* n 3)) 2)))]))
  
  ;; collatz/hash: number -> number.
  (define/memo (collatz/hash n) memoize/hash
    (cond [(= n 1) 1]
          [(even? n)
           (add1 (collatz/hash (/ n 2)))]
          [else
           (+ 2 (collatz/hash (/ (+ 1 (* n 3)) 2)))]))
  
  
  ;; foldr-range: (number X -> X) X number number -> X
  ;; Folds a function across the half-open interval [a, b),
  ;; starting from b and counting down.
  (define (foldr-range f acc a b)
    (let loop ([acc acc]
               [i (sub1 b)])
      (cond [(<= a i)
             (loop (f i acc) (sub1 i))]
            [else acc])))
  
  
  ;; find-max-chain: number number -> number
  ;; Returns the maximum chain length for chains whose
  ;; first number is between i and j, inclusive.
  (define (find-max-chain i j f)
    (local ((define (update-max x best-so-far)
              (max (f x) best-so-far)))
      (foldr-range update-max 0 i (add1 j))))
  
  
  (begin (printf "vector: ")
         (collect-garbage)
         (time (find-max-chain 1 1000000 collatz)))
  (begin (printf "hash: ")
         (collect-garbage)
         (time (find-max-chain 1 1000000 collatz/hash))))