blob: 768aee28e3620abb4ca9b4724e36faeffb28ffe4 (
about) (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
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)))
|
