; google code jam qualification round africa 2010, revisited
(define (make-dict lt?)
(define-syntax define-generator
(lambda (x)
(syntax-case x (lambda)
((stx name (lambda formals e0 e1 ...))
(with-syntax ((yield (datum->syntax-object (syntax stx) 'yield)))
(syntax (define name
(lambda formals
(let ((resume #f) (return #f))
(define yield
(lambda args
(call-with-current-continuation
(lambda (cont)
(set! resume cont)
(apply return args)))))
(lambda ()
(call-with-current-continuation
(lambda (cont)
(set! return cont)
(cond (resume (resume))
(else (let () e0 e1 ...)
(error 'name "unexpected return"))))))))))))
((stx (name . formals) e0 e1 ...)
(syntax (stx name (lambda formals e0 e1 ...)))))))
(define (tree k v l r)
(vector k v l r (+ (max (ht l) (ht r)) 1)
(+ (size l) (size r) 1)))
(define (key t) (vector-ref t 0))
(define (val t) (vector-ref t 1))
(define (lkid t) (vector-ref t 2))
(define (rkid t) (vector-ref t 3))
(define (ht t) (vector-ref t 4))
(define (size t) (vector-ref t 5))
(define (bal t) (- (ht (lkid t)) (ht (rkid t))))
(define nil (vector 'nil 'nil 'nil 'nil 0 0))
(define (nil? t) (eq? t nil))
(define (rot-left t)
(if (nil? t) t
(tree (key (rkid t))
(val (rkid t))
(tree (key t) (val t) (lkid t) (lkid (rkid t)))
(rkid (rkid t)))))
(define (rot-right t)
(if (nil? t) t
(tree (key (lkid t))
(val (lkid t))
(lkid (lkid t))
(tree (key t) (val t) (rkid (lkid t)) (rkid t)))))
(define (balance t)
(let ((b (bal t)))
(cond ((< (abs b) 2) t)
((positive? b)
(if (< -1 (bal (lkid t))) (rot-right t)
(rot-right (tree (key t) (val t)
(rot-left (lkid t)) (rkid t)))))
((negative? b)
(if (< (bal (rkid t)) 1) (rot-left t)
(rot-left (tree (key t) (val t)
(lkid t) (rot-right (rkid t)))))))))
(define (lookup t k)
(cond ((nil? t) #f)
((lt? k (key t)) (lookup (lkid t) k))
((lt? (key t) k) (lookup (rkid t) k))
(else (cons k (val t)))))
(define (insert t k v)
(cond ((nil? t) (tree k v nil nil))
((lt? k (key t))
(balance (tree (key t) (val t)
(insert (lkid t) k v) (rkid t))))
((lt? (key t) k)
(balance (tree (key t) (val t)
(lkid t) (insert (rkid t) k v))))
(else (tree k v (lkid t) (rkid t)))))
(define (update t f k v)
(cond ((nil? t) (tree k v nil nil))
((lt? k (key t))
(balance (tree (key t) (val t)
(update (lkid t) f k v) (rkid t))))
((lt? (key t) k)
(balance (tree (key t) (val t)
(lkid t) (update (rkid t) f k v))))
(else (tree k (f k (val t)) (lkid t) (rkid t)))))
(define (delete-successor t)
(if (nil? (lkid t)) (values (rkid t) (key t) (val t))
(call-with-values
(lambda () (delete-successor (lkid t)))
(lambda (l k v)
(values (balance (tree (key t) (val t) l (rkid t))) k v)))))
(define (delete t k)
(cond ((nil? t) nil)
((lt? k (key t))
(balance (tree (key t) (val t)
(delete (lkid t) k) (rkid t))))
((lt? (key t) k)
(balance (tree (key t) (val t)
(lkid t) (delete (rkid t) k))))
((nil? (lkid t)) (rkid t))
((nil? (rkid t)) (lkid t))
(else (call-with-values
(lambda () (delete-successor (rkid t)))
(lambda (r k v) (balance (tree k v (lkid t) r)))))))
(define (nth t n)
(if (negative? n) (error 'nth "must be non-negative")
(let ((s (size (lkid t))))
(cond ((< n s) (nth (lkid t) n))
((< s n) (nth (rkid t) (- n s 1)))
((nil? t) #f)
(else (cons (key t) (val t)))))))
(define (rank t k)
(let loop ((t t) (s (size (lkid t))))
(cond ((nil? t) #f)
((lt? k (key t))
(loop (lkid t) (size (lkid (lkid t)))))
((lt? (key t) k)
(loop (rkid t) (+ s (size (lkid (rkid t))) 1)))
(else s))))
(define (avl-map proc t) ; (proc key value)
(if (nil? t) nil
(tree (key t) (proc (key t) (val t))
(avl-map proc (lkid t))
(avl-map proc (rkid t)))))
(define (avl-fold proc base t) ; (proc key value base)
(if (nil? t) base
(avl-fold proc
(proc (key t) (val t)
(avl-fold proc base (lkid t)))
(rkid t))))
(define (avl-for-each proc t) ; (proc key value)
(unless (nil? t)
(avl-for-each proc (lkid t))
(proc (key t) (val t))
(avl-for-each proc (rkid t))))
(define (to-list t)
(cond ((nil? t) (list))
((and (nil? (lkid t)) (nil? (rkid t)))
(list (cons (key t) (val t))))
(else (append (to-list (lkid t))
(list (cons (key t) (val t)))
(to-list (rkid t))))))
(define (from-list t xs)
(let loop ((xs xs) (t t))
(if (null? xs) t
(loop (cdr xs) (insert t (caar xs) (cdar xs))))))
(define-generator (make-gen t)
(avl-for-each (lambda (k v) (yield (cons k v))) t)
(do () (#f) (yield #f)))
(define (new dict)
(lambda (message . args) (dispatch dict message args)))
(define (dispatch dict message args)
(define (arity n)
(if (not (= (length args) n)) (error 'dict "incorrect arity")))
(case message
((empty? nil?) (arity 0) (nil? dict))
((lookup fetch get) (arity 1) (apply lookup dict args))
((insert store put) (arity 2) (new (apply insert dict args)))
((update) (arity 3) (new (apply update dict args)))
((delete remove) (arity 1) (new (apply delete dict args)))
((size count length) (arity 0) (size dict))
((nth) (arity 1) (apply nth dict args))
((rank) (arity 1) (apply rank dict args))
((map) (arity 1) (new (avl-map (car args) dict)))
((fold) (arity 2) (avl-fold (car args) (cadr args) dict))
((for-each) (arity 1) (avl-for-each (car args) dict))
((to-list enlist) (arity 0) (to-list dict))
((from-list) (arity 1) (new (apply from-list dict args)))
((make-gen gen) (arity 0) (make-gen dict))
(else (error 'dict "invalid message"))))
(vector-set! nil 2 nil) (vector-set! nil 3 nil) (new nil))
(define (make-hash hash eql? oops size)
(let ((table (make-vector size '())))
(lambda (message . args)
(if (eq? message 'enlist)
(let loop ((k 0) (result '()))
(if (= size k)
result
(loop (+ k 1) (append (vector-ref table k) result))))
(let* ((key (car args))
(index (modulo (hash key) size))
(bucket (vector-ref table index)))
(case message
((lookup fetch get ref recall)
(let loop ((bucket bucket))
(cond ((null? bucket) oops)
((eql? (caar bucket) key) (cdar bucket))
(else (loop (cdr bucket))))))
((insert insert! ins ins! set set! store store! install install!)
(vector-set! table index
(let loop ((bucket bucket))
(cond ((null? bucket)
(list (cons key (cadr args))))
((eql? (caar bucket) key)
(cons (cons key (cadr args)) (cdr bucket)))
(else (cons (car bucket) (loop (cdr bucket))))))))
((delete delete! del del! remove remove!)
(vector-set! table index
(let loop ((bucket bucket))
(cond ((null? bucket) '())
((eql? (caar bucket) key)
(cdr bucket))
(else (cons (car bucket) (loop (cdr bucket))))))))
((update update!)
(vector-set! table index
(let loop ((bucket bucket))
(cond ((null? bucket)
(list (cons key (caddr args))))
((eql? (caar bucket) key)
(cons (cons key ((cadr args) key (cdar bucket))) (cdr bucket)))
(else (cons (car bucket) (loop (cdr bucket))))))))
(else (error 'hash-table "unrecognized message")) ))))))
(define sort #f)
(define merge #f)
(let ()
(define dosort
(lambda (pred? ls n)
(if (= n 1)
(list (car ls))
(let ((i (quotient n 2)))
(domerge pred?
(dosort pred? ls i)
(dosort pred? (list-tail ls i) (- n i)))))))
(define domerge
(lambda (pred? l1 l2)
(cond
((null? l1) l2)
((null? l2) l1)
((pred? (car l2) (car l1))
(cons (car l2) (domerge pred? l1 (cdr l2))))
(else (cons (car l1) (domerge pred? (cdr l1) l2))))))
(set! sort
(lambda (pred? l)
(if (null? l) l (dosort pred? l (length l)))))
(set! merge
(lambda (pred? l1 l2)
(domerge pred? l1 l2))))
(define (vector-sort! vec comp)
(define-syntax while
(syntax-rules ()
((while pred? body ...)
(do () ((not pred?)) body ...))))
(define-syntax assign!
(syntax-rules ()
((assign! var expr)
(begin (set! var expr) var))))
(define len (vector-length vec))
(define-syntax v (syntax-rules () ((v k) (vector-ref vec k))))
(define-syntax v! (syntax-rules () ((v! k x) (vector-set! vec k x))))
(define-syntax cmp (syntax-rules () ((cmp a b) (comp (v a) (v b)))))
(define-syntax lt? (syntax-rules () ((lt? a b) (negative? (cmp a b)))))
(define-syntax swap! (syntax-rules () ((swap! a b)
(let ((t (v a))) (v! a (v b)) (v! b t)))))
(define (vecswap! a b s)
(do ((a a (+ a 1)) (b b (+ b 1)) (s s (- s 1))) ((zero? s))
(swap! a b)))
(define (med3 a b c)
(if (lt? b c)
(if (lt? b a) (if (lt? c a) c a) b)
(if (lt? c a) (if (lt? b a) b a) c)))
(define (pv-init a n)
(let ((pm (+ a (quotient n 2))))
(when (> n 7)
(let ((pl a) (pn (+ a n -1)))
(when (> n 40)
(let ((s (quotient n 8)))
(set! pl (med3 pl (+ pl s) (+ pl s s)))
(set! pm (med3 (- pm s) pm (+ pm s)))
(set! pn (med3 (- pn s s) (- pn s) pn))))
(set! pm (med3 pl pm pn))))
pm))
(let qsort ((a 0) (n len))
(if (< n 7)
(do ((pm (+ a 1) (+ pm 1))) ((not (< pm (+ a n))))
(do ((pl pm (- pl 1)))
((not (and (> pl a) (> (cmp (- pl 1) pl) 0))))
(swap! pl (- pl 1))))
(let ((pv (pv-init a n)) (r #f)
(pa a) (pb a) (pc (+ a n -1)) (pd (+ a n -1)))
(swap! a pv) (set! pv a)
(let loop ()
(while (and (<= pb pc) (<= (assign! r (cmp pb pv)) 0))
(when (= r 0) (swap! pa pb) (set! pa (+ pa 1)))
(set! pb (+ pb 1)))
(while (and (>= pc pb) (>= (assign! r (cmp pc pv)) 0))
(when (= r 0) (swap! pc pd) (set! pd (- pd 1)))
(set! pc (- pc 1)))
(unless (> pb pc)
(swap! pb pc) (set! pb (+ pb 1)) (set! pc (- pc 1)) (loop)))
(let ((pn (+ a n)))
(let ((s (min (- pa a) (- pb pa)))) (vecswap! a (- pb s) s))
(let ((s (min (- pd pc) (- pn pd 1)))) (vecswap! pb (- pn s) s))
(let ((s (- pb pa))) (when (> s 1) (qsort a s)))
(let ((s (- pd pc))) (when (> s 1) (qsort (- pn s) s))))))))
(define (identity x) x)
(define (soln1 target vec)
(let ((len (vector-length vec)))
(call-with-current-continuation
(lambda (return)
(do ((i 0 (+ i 1))) ((= i len) #f)
(do ((j (+ i 1) (+ j 1))) ((= j len))
(when (= target (+ (vector-ref vec i) (vector-ref vec j)))
(return (list (+ i 1) (+ j 1))))))))))
(define (soln2 target xs)
(let ((items (make-dict <)))
(do ((i 0 (+ i 1)) (xs xs (cdr xs))) ((null? xs))
(set! items (items 'insert (car xs) (+ i 1))))
(let loop ((xs xs) (i 1))
(cond ((null? xs) #f)
((items 'lookup (- target (car xs))) =>
(lambda (item)
(if (= i (cdr item))
(loop (cdr xs) (+ i 1))
(sort < (list i (cdr item))))))
(else (loop (cdr xs) (+ i 1)))))))
(define (soln3 target xs)
(let ((items (make-hash identity = #f 97)))
(do ((i 0 (+ i 1)) (xs xs (cdr xs))) ((null? xs))
(items 'insert (car xs) (+ i 1)))
(let loop ((xs xs) (i 1))
(cond ((null? xs) #f)
((items 'lookup (- target (car xs))) =>
(lambda (item)
(if (= i item)
(loop (cdr xs) (+ i 1))
(sort < (list i item)))))
(else (loop (cdr xs) (+ i 1)))))))
(define (soln4 target xs)
(define (comp a b)
(if (< (car a) (car b)) -1
(if (< (car b) (car a)) 1 0)))
(define (bsearch target vec lo hi)
(if (< hi lo) #f
(let ((mid (quotient (+ lo hi) 2)))
(cond ((< (car (vector-ref vec mid)) target)
(bsearch target vec (+ mid 1) hi))
((< target (car (vector-ref vec mid)))
(bsearch target vec lo (- mid 1)))
(else (vector-ref vec mid))))))
(let loop ((xs xs) (i 0) (vs (list)))
(if (pair? xs)
(loop (cdr xs) (+ i 1) (cons (cons (car xs) i) vs))
(let ((vec (list->vector vs)))
(vector-sort! vec comp)
(let ((len (vector-length vec)))
(let loop ((i 0))
(cond ((= i len) #f)
((bsearch (- target (car (vector-ref vec i)))
vec (+ i 1) (- len 1)) =>
(lambda (item)
(sort < (list (+ i 2) (+ (cdr item) 1)))))
(else (loop (+ i 1))))))))))
(display (soln1 100 '#(5 75 25))) (newline)
(display (soln1 200 '#(150 24 79 50 88 345 3))) (newline)
(display (soln1 8 '#(2 1 9 4 4 56 90 3))) (newline)
(display (soln1 17 '#(2 1 9 4 4 56 90 3))) (newline)
(newline)
(display (soln2 100 '(5 75 25))) (newline)
(display (soln2 200 '(150 24 79 50 88 345 3))) (newline)
(display (soln2 8 '(2 1 9 4 4 56 90 3))) (newline)
(display (soln2 17 '(2 1 9 4 4 56 90 3))) (newline)
(newline)
(display (soln3 100 '(5 75 25))) (newline)
(display (soln3 200 '(150 24 79 50 88 345 3))) (newline)
(display (soln3 8 '(2 1 9 4 4 56 90 3))) (newline)
(display (soln3 17 '(2 1 9 4 4 56 90 3))) (newline)
(newline)
(display (soln4 100 '(5 75 25))) (newline)
(display (soln4 200 '(150 24 79 50 88 345 3))) (newline)
(display (soln4 8 '(2 1 9 4 4 56 90 3))) (newline)
(display (soln4 17 '(2 1 9 4 4 56 90 3))) (newline)