From 1feb32ab54b5955464db80919d24716a8ba477c1 Mon Sep 17 00:00:00 2001 From: Stefan Israelsson Tampe Date: Wed, 28 Feb 2018 13:24:08 +0100 Subject: heapq added --- modules/language/python/module/heapq.scm | 236 +++++++++++++++++++++++++++++++ 1 file changed, 236 insertions(+) create mode 100644 modules/language/python/module/heapq.scm (limited to 'modules/language/python/module/heapq.scm') diff --git a/modules/language/python/module/heapq.scm b/modules/language/python/module/heapq.scm new file mode 100644 index 0000000..768aee2 --- /dev/null +++ b/modules/language/python/module/heapq.scm @@ -0,0 +1,236 @@ +(define-module (language python module heapq) + #:use-module (ice-9 control) + #:use-module (language python for) + #:use-module (language python yield) + #:use-module (language python try) + #:use-module (language python exceptions) + #:use-module (language python list) + #:use-module (language python def) + #:use-module (oop pf-objects) + #:use-module ((language python module python) + #:select (zip sorted range enumerate iter max min)) + #:export (heappush heappop heapify heapreplace merge nlargest nsmallest + heappushpop)) + +(define-syntax-rule (aif it p a b) (let ((it p)) (if it a b))) + +(define (heappush heap item) + (pylist-append! heap item) + (siftdown heap 0 (- (len heap) 1))) + +(define (heappop heap) + (let ((lastelt (pylist-pop! heap))) + (if (> (len heap) 0) + (let ((ret (pylist-ref heap 0))) + (pylist-set! heap 0 lastelt) + (siftup heap 0) + ret) + lastelt))) + +(define (heapreplace heap item) + (let ((ret (pylist-ref heap 0))) + (pylist-set! heap 0 item) + (siftup heap 0) + ret)) + +(define (heapify x) + (let lp ((i (- (floor-quotient (len x) 2) 1))) + (if (>= i 0) + (begin + (siftup x i) + (lp (- i 1)))))) + +(define (heappop_max heap) + (let ((lastelt (pylist-pop! heap))) + (if (> (len heap) 0) + (let ((ret (pylist-ref heap 0))) + (pylist-set! heap 0 lastelt) + (siftup_max heap 0) + ret) + lastelt))) + +(define (heapreplace_max heap item) + (let ((ret (pylist-ref heap 0))) + (pylist-set! heap 0 item) + (siftup_max heap 0) + ret)) + +(define (heapify_max x) + (let lp ((i (- (floor-quotient (len x) 2) 1))) + (if (>= i 0) + (begin + (siftup_max x i) + (lp (- i 1)))))) + +(define (siftdown heap startpos pos) + (define newitem (pylist-ref heap pos)) + + (let lp ((pos pos)) + (if (> pos startpos) + (let () + (define parentpos (ash (- pos 1) 1)) + (define parent (pylist-ref heap parentpos)) + (if (< newitem parent) + (begin + (pylist-set! heap pos parent) + (lp parentpos)) + (pylist-set! heap pos newitem))) + (pylist-set! heap pos newitem)))) + +(define (siftup heap pos) + (define endpos (len heap)) + (define startpos pos) + (define newitem (pylist-ref heap pos)) + (let lp ((pos pos) (childpos (+ (* 2 pos) 1))) + (if (< childpos endpos) + (let* ((rightpos (+ childpos 1)) + (childpos + (if (and (< rightpos endpos) + (not (< (pylist-ref heap childpos) + (pylist-ref heap rightpos)))) + rightpos + childpos))) + (pylist-set! heap pos (pylist-ref heap childpos)) + (lp childpos (+ (* 2 childpos) 1))) + (begin + (pylist-set! heap pos newitem) + (siftdown heap startpos pos))))) + +(define (siftdown_max heap startpos pos) + (define newitem (pylist-ref heap pos)) + + (let lp ((pos pos)) + (if (> pos startpos) + (let () + (define parentpos (ash (- pos 1) 1)) + (define parent (pylist-ref heap parentpos)) + (if (< parent newitem) + (begin + (pylist-set! heap pos parent) + (lp parentpos)) + (pylist-set! heap pos newitem))) + (pylist-set! heap pos newitem)))) + +(define (siftup_max heap pos) + (define endpos (len heap)) + (define startpos pos) + (define newitem (pylist-ref heap pos)) + (let lp ((pos pos) (childpos (+ (* 2 pos) 1))) + (if (< childpos endpos) + (let* ((rightpos (+ childpos 1)) + (childpos + (if (and (< rightpos endpos) + (not (< (pylist-ref heap rightpos) + (pylist-ref heap childpos)))) + rightpos + childpos))) + (pylist-set! heap pos (pylist-ref heap childpos)) + (lp childpos (+ (* 2 childpos) 1))) + (begin + (pylist-set! heap pos newitem) + (siftdown_max heap startpos pos))))) + +(def (merge (* iterables) (= key (lambda (x) x)) (= reverse #f)) + ((make-generator () + (lambda (yield) + (define h (py-list)) + (define (h_append x) (pylist-append! h x)) + (define-values (-heapify -heappop -heapreplace direction) + (if reverse + (values heapify_max + heappop_max + heapreplace_max + -1) + (values heapify + heappop + heapreplace + 1))) + + (for ((order it : (enumerate (map iter iterables)))) () + (define next- (aif x (ref it '__next__) + x + (lambda () (next it)))) + (try + (lambda () + (h_append (vector (key (next-)) + (* order direction) + next-))) + (#:except StopIteration => + (lambda x (-heappop h))))) + + (-heapify h) + (let lp () + (if (> (len h) 1) + (try + (lambda () + (let lp () + (let ((s (pylist-ref h 0))) + (yield (vector-ref s 0)) + (vector-set! s 0 (key ((vector-ref s 2)))) + (-heapreplace h s) + (lp)))) + (#:except StopIteration => + (lambda x (-heappop h) (lp)))))) + + (if (> (len h) 0) + (let* ((x (pylist-ref h 0)) + (n (pylist-ref x 2))) + (let lp ((i (vector-ref x 0))) + (yield i) + (lp (n))))))))) + + + + +(define-syntax-rule (mkn nsmallest heapify_max heapreplace_max q< ran o ++) + (define* (nsmallest n iterable #:key (key (lambda (x) x))) + (let/ec return + (if (equal? n 1) + (return + (let* ((sent (list 1)) + (res (min (iter iterable) #:default sent #:key key))) + (if (eq? sent res) + '() + (list res))))) + + (let ((size (try + (lambda () (len iterable)) + (#:except (or TypeError AttributeError) => + (lambda x #f))))) + (if (and size (>= n size)) + (return (pylist-slice (sorted iterable #:key key) 0 + ((@ (guile) min) n size) + 1)))) + + (let ((it (iter iterable)) + (result (py-list))) + + (for ((i elem : (zip (ran n) it))) () + (pylist-append! result (vector (key elem) i elem))) + + (if (= 0 (len result)) + (return result)) + + (heapify_max result) + + (for ((elem : it)) ((top (vector-ref (pylist-ref result 0) 0)) + (order (o n))) + (let ((k (key elem))) + (if (q< k top) + (begin + (heapreplace_max result (vector k order elem)) + (values (vector-ref (pylist-ref result 0) 0) + (++ order))) + (values top order)))) + (pylist-sort! result) + (return + (let ((r (py-list))) + (for ((x : result)) () + (pylist-append! r (vector-ref x 2))) + r)))))) + +(mkn nsmallest heapify_max heapreplace_max < (lambda (n) (range n)) + (lambda (n) n) (lambda (x) (+ x 1))) + +(mkn nlargest heapify heapreplace > (lambda (n) (range 0 (- n) -1)) + (lambda (n) (- n)) (lambda (x) (- x 1))) -- cgit v1.2.3