heap.bend
bend-kit-collections@0.1.2.0 · 130 lines · 0 definitions
1# Heap: a priority queue, as a skew heap ordered by a le you pass in.
2#
3# The smallest element (by le) sits at the top. push and pop are
4# amortized O(log n) over a heap used once (sharing one with + and popping
5# both copies can cost more), peek and size are O(1), and sort is
6# heapsort. le is a template, as in List.sort: Heap.push(~U32, ~U32.is_le,
7# 7, h).
8#
9# The laws at the bottom are checked every time this file is imported.
10import Base
11
12type Heap<-A: Data> is Data:
13 HLeaf{}
14 HNode{size: Nat, top: A, left: Heap<A>, right: Heap<A>}
15
16def Heap.new(-A: Data) -> Heap<A>:
17 HLeaf{}
18
19def Heap.size(-A: Data, h: Heap<A>) -> Nat:
20 match h:
21 case HLeaf{}:
22 0n
23 case HNode{n, x, l, r}:
24 n
25
26def Heap.peek(-A: Data, h: Heap<A>) -> Maybe<&2, A>:
27 match h:
28 case HLeaf{}:
29 None{}
30 case HNode{n, x, l, r}:
31 Some{x}
32
33# The smaller root wins: its right child merges with the other heap and
34# becomes the left, and its old left becomes the right (a skew heap's
35# swap, which keeps merges short on average). The fuel is the size of both
36# heaps: every step takes a root out, so it never runs short.
37def Heap.merge.go(
38 ~A: Data, ~le: A -> A -> Bool, fuel: Nat, a: Heap<A>, b: Heap<A>
39) -> Heap<A>:
40 match fuel:
41 case 0n:
42 a
43 case 1n+f:
44 match a b:
45 case HLeaf{} HLeaf{}:
46 HLeaf{}
47 case HLeaf{} HNode{nb, y, bl, br}:
48 HNode{nb, y, bl, br}
49 case HNode{na, x, al, ar} HLeaf{}:
50 HNode{na, x, al, ar}
51 case HNode{+na, +x, +al, +ar} HNode{+nb, +y, +bl, +br}:
52 +c = le(x, y)
53 HNode{Nat.add(na, nb), Bool.pick(A, c, x, y),
54 Heap.merge.go(~A, ~le, f,
55 Bool.pick(Heap<A>, c, ar, br),
56 Bool.pick(Heap<A>, c, HNode{nb, y, bl, br},
57 HNode{na, x, al, ar})),
58 Bool.pick(Heap<A>, c, al, bl)}
59
60def Heap.merge(
61 ~A: Data, ~le: A -> A -> Bool, +a: Heap<A>, +b: Heap<A>
62) -> Heap<A>:
63 Heap.merge.go(~A, ~le, Nat.add(Heap.size(A, a), Heap.size(A, b)), a, b)
64
65def Heap.push(~A: Data, ~le: A -> A -> Bool, x: A, h: Heap<A>) -> Heap<A>:
66 Heap.merge(~A, ~le, HNode{1n, x, HLeaf{}, HLeaf{}}, h)
67
68# the top, and the heap without it; None on an empty heap
69def Heap.pop(
70 ~A: Data, ~le: A -> A -> Bool, h: Heap<A>
71) -> Maybe<&1, A & Heap<A>>:
72 match h:
73 case HLeaf{}:
74 None{}
75 case HNode{n, x, l, r}:
76 Some{(x, Heap.merge(~A, ~le, l, r))}
77
78def Heap.from_list(~A: Data, ~le: A -> A -> Bool, xs: List<&2, A>) -> Heap<A>:
79 match xs:
80 case Nil{}:
81 HLeaf{}
82 case h <> t:
83 Heap.push(~A, ~le, h, Heap.from_list(~A, ~le, t))
84
85def Heap.to_list.go(
86 ~A: Data, ~le: A -> A -> Bool, fuel: Nat, h: Heap<A>
87) -> List<&2, A>:
88 match fuel:
89 case 0n:
90 Nil{}
91 case 1n+f:
92 match h:
93 case HLeaf{}:
94 Nil{}
95 case HNode{n, x, l, r}:
96 x <> Heap.to_list.go(~A, ~le, f, Heap.merge(~A, ~le, l, r))
97
98# the elements, smallest first
99def Heap.to_list(~A: Data, ~le: A -> A -> Bool, +h: Heap<A>) -> List<&2, A>:
100 Heap.to_list.go(~A, ~le, Heap.size(A, h), h)
101
102# heapsort
103def Heap.sort(~A: Data, ~le: A -> A -> Bool, xs: List<&2, A>) -> List<&2, A>:
104 Heap.to_list(~A, ~le, Heap.from_list(~A, ~le, xs))
105
106# Laws
107# ----
108
109# LAW: popping what was pushed onto an empty heap gives it back
110law Heap.pop.push_new:
111 for -x: U32
112 {Heap.pop(~U32, ~U32.is_le, Heap.push(~U32, ~U32.is_le, x, Heap.new(U32)))
113 == Some{(x, Heap.new(U32))} : Maybe<&1, U32 & Heap<U32>>}
114
115def Heap.pop.push_new(x):
116 {==}
117
118# LAW: a push adds one to the size
119law Heap.size.push:
120 for -x: U32
121 for h: Heap<U32>
122 {Heap.size(U32, Heap.push(~U32, ~U32.is_le, x, h))
123 == 1n+Heap.size(U32, h) : Nat}
124
125def Heap.size.push(x, h):
126 match h:
127 case HLeaf{}:
128 {==}
129 case HNode{n, y, l, r}:
130 {==}