diff --git a/stdlib/src/heapq/__init__.mojo b/stdlib/src/heapq/__init__.mojo new file mode 100644 index 00000000000..5d5616bdeb7 --- /dev/null +++ b/stdlib/src/heapq/__init__.mojo @@ -0,0 +1,14 @@ +# ===----------------------------------------------------------------------=== # +# Copyright (c) 2024, Modular Inc. All rights reserved. +# +# Licensed under the Apache License v2.0 with LLVM Exceptions: +# https://llvm.org/LICENSE.txt +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ===----------------------------------------------------------------------=== # + +from .heapq import heapify, heappush, heappop diff --git a/stdlib/src/heapq/heapq.mojo b/stdlib/src/heapq/heapq.mojo new file mode 100644 index 00000000000..ff6b756eb10 --- /dev/null +++ b/stdlib/src/heapq/heapq.mojo @@ -0,0 +1,95 @@ +# ===----------------------------------------------------------------------=== # +# Copyright (c) 2024, Modular Inc. All rights reserved. +# +# Licensed under the Apache License v2.0 with LLVM Exceptions: +# https://llvm.org/LICENSE.txt +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ===----------------------------------------------------------------------=== # +"""Defines the heapq module. + +Implementation currently tightly follows the Python implementation. +""" + + +fn heappush[T: ComparableCollectionElement](inout heap: List[T], owned item: T): + """Push an element onto a heapified list. + + Parameters: + T: A comparable collection element type. + + Args: + heap: The heap to push to. + item: The new item to be placed in the heap. + """ + heap.append(item^) + _siftdown(heap, 0, len(heap) - 1) + + +fn heappop[T: ComparableCollectionElement](inout heap: List[T]) -> T: + """Pop an element from a heapified list. + + Parameters: + T: A comparable collection element type. + + Args: + heap: The heap to push to. + + Returns: + The popped element. + """ + var lastelt = heap.pop() + if heap: + var returnitem = heap[0] + heap[0] = lastelt + _siftup(heap, 0) + return returnitem + return lastelt + + +fn heapify[T: ComparableCollectionElement](inout x: List[T]): + """Convert a list of elements into a binary heap. + + Parameters: + T: A comparable collection element type. + + Args: + x: The list to heapify. + """ + for i in reversed(range(len(x) // 2)): + _siftup(x, i) + + +fn _siftdown[ + T: ComparableCollectionElement +](inout heap: List[T], startpos: Int, owned pos: Int): + var newitem = heap[pos] + while pos > startpos: + var parentpos = (pos - 1) >> 1 + var parent = heap[parentpos] + if newitem < parent: + heap[pos] = parent + pos = parentpos + continue + break + heap[pos] = newitem + + +fn _siftup[T: ComparableCollectionElement](inout heap: List[T], owned pos: Int): + var endpos = len(heap) + var startpos = pos + var newitem = heap[pos] + var childpos = 2 * pos + 1 + while childpos < endpos: + var rightpos = childpos + 1 + if rightpos < endpos and not heap[childpos] < heap[rightpos]: + childpos = rightpos + heap[pos] = heap[childpos] + pos = childpos + childpos = 2 * pos + 1 + heap[pos] = newitem + _siftdown(heap, startpos, pos) diff --git a/stdlib/test/heapq/test_heapq.mojo b/stdlib/test/heapq/test_heapq.mojo new file mode 100644 index 00000000000..38aef866836 --- /dev/null +++ b/stdlib/test/heapq/test_heapq.mojo @@ -0,0 +1,36 @@ +# ===----------------------------------------------------------------------=== # +# Copyright (c) 2024, Modular Inc. All rights reserved. +# +# Licensed under the Apache License v2.0 with LLVM Exceptions: +# https://llvm.org/LICENSE.txt +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ===----------------------------------------------------------------------=== # +# RUN: %mojo %s + +from heapq import heapify, heappush, heappop +from testing import assert_equal + + +def main(): + var l = List(3, 6, 7, 8, 3, 7) + heapify(l) + # TODO: use __eq__ when implemented for List + assert_equal(l.__str__(), "[3, 3, 7, 8, 6, 7]") + heappush(l, 5) + assert_equal(l.__str__(), "[3, 3, 5, 8, 6, 7, 7]") + + assert_equal(heappop(l), 3) + assert_equal(l.__str__(), "[3, 6, 5, 8, 7, 7]") + + l = List(57, 3467, 734, 4, 6, 8, 236, 367, 236, 75, 87) + heapify(l) + assert_equal(l.__str__(), "[4, 6, 8, 236, 57, 734, 236, 367, 3467, 75, 87]") + + # sanity check that nothing bad happens + var l2 = List[Int]() + heapify(l2)