heap.bend

bend-kit-collections@0.1.4.1 · 130 lines · 14 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      {==}