Heap Implementation in Python

heap.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "left = lambda x: x * 2 + 1\n",
    "right = lambda x: x * 2 + 2\n",
    "parent = lambda x: (x - 1) // 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def push(heap, element):\n",
    "    i = len(heap)\n",
    "    heap.append(element)\n",
    "    while i:\n",
    "        p = parent(i)\n",
    "        if heap[i] >= heap[p]:\n",
    "            break\n",
    "        \n",
    "        heap[i], heap[p] = heap[p], heap[i]\n",
    "        i = p\n",
    "\n",
    "def pop(heap):\n",
    "    if len(heap) == 1:\n",
    "        return heap.pop()\n",
    "    \n",
    "    out = heap[0]\n",
    "    \n",
    "    heap[0] = heap.pop()\n",
    "    i = 0\n",
    "    while (l := left(i)) < len(heap):\n",
    "        k = l if (r := right(i)) >= len(heap) or heap[l] < heap[r] else r\n",
    "        \n",
    "        if heap[i] <= heap[k]:\n",
    "            break\n",
    "        \n",
    "        heap[k], heap[i] = heap[i], heap[k]\n",
    "        i = k\n",
    "    \n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "push 11 -> [11]\n",
      "push 20 -> [11, 20]\n",
      "push 19 -> [11, 20, 19]\n",
      "push 3 -> [3, 11, 19, 20]\n",
      "push 15 -> [3, 11, 19, 20, 15]\n",
      "push 4 -> [3, 11, 4, 20, 15, 19]\n",
      "push 10 -> [3, 11, 4, 20, 15, 19, 10]\n",
      "push 8 -> [3, 8, 4, 11, 15, 19, 10, 20]\n",
      "push 8 -> [3, 8, 4, 8, 15, 19, 10, 20, 11]\n",
      "push 7 -> [3, 7, 4, 8, 8, 19, 10, 20, 11, 15]\n",
      "push 11 -> [3, 7, 4, 8, 8, 19, 10, 20, 11, 15, 11]\n",
      "push 17 -> [3, 7, 4, 8, 8, 17, 10, 20, 11, 15, 11, 19]\n",
      "push 17 -> [3, 7, 4, 8, 8, 17, 10, 20, 11, 15, 11, 19, 17]\n",
      "push 6 -> [3, 7, 4, 8, 8, 17, 6, 20, 11, 15, 11, 19, 17, 10]\n",
      "push 8 -> [3, 7, 4, 8, 8, 17, 6, 20, 11, 15, 11, 19, 17, 10, 8]\n",
      "push 1 -> [1, 3, 4, 7, 8, 17, 6, 8, 11, 15, 11, 19, 17, 10, 8, 20]\n",
      "push 1 -> [1, 1, 4, 3, 8, 17, 6, 7, 11, 15, 11, 19, 17, 10, 8, 20, 8]\n",
      "push 14 -> [1, 1, 4, 3, 8, 17, 6, 7, 11, 15, 11, 19, 17, 10, 8, 20, 8, 14]\n",
      "push 3 -> [1, 1, 4, 3, 8, 17, 6, 7, 3, 15, 11, 19, 17, 10, 8, 20, 8, 14, 11]\n",
      "push 18 -> [1, 1, 4, 3, 8, 17, 6, 7, 3, 15, 11, 19, 17, 10, 8, 20, 8, 14, 11, 18]\n",
      "push 6 -> [1, 1, 4, 3, 6, 17, 6, 7, 3, 8, 11, 19, 17, 10, 8, 20, 8, 14, 11, 18, 15]\n",
      "push 7 -> [1, 1, 4, 3, 6, 17, 6, 7, 3, 8, 7, 19, 17, 10, 8, 20, 8, 14, 11, 18, 15, 11]\n",
      "push 4 -> [1, 1, 4, 3, 4, 17, 6, 7, 3, 8, 6, 19, 17, 10, 8, 20, 8, 14, 11, 18, 15, 11, 7]\n",
      "push 1 -> [1, 1, 1, 3, 4, 4, 6, 7, 3, 8, 6, 17, 17, 10, 8, 20, 8, 14, 11, 18, 15, 11, 7, 19]\n",
      "push 16 -> [1, 1, 1, 3, 4, 4, 6, 7, 3, 8, 6, 16, 17, 10, 8, 20, 8, 14, 11, 18, 15, 11, 7, 19, 17]\n",
      "pop -> 1 [1, 1, 4, 3, 4, 16, 6, 7, 3, 8, 6, 17, 17, 10, 8, 20, 8, 14, 11, 18, 15, 11, 7, 19]\n",
      "pop -> 1 [1, 3, 4, 3, 4, 16, 6, 7, 11, 8, 6, 17, 17, 10, 8, 20, 8, 14, 19, 18, 15, 11, 7]\n",
      "pop -> 1 [3, 3, 4, 7, 4, 16, 6, 7, 11, 8, 6, 17, 17, 10, 8, 20, 8, 14, 19, 18, 15, 11]\n",
      "pop -> 3 [3, 4, 4, 7, 6, 16, 6, 7, 11, 8, 11, 17, 17, 10, 8, 20, 8, 14, 19, 18, 15]\n",
      "pop -> 3 [4, 4, 6, 7, 6, 16, 8, 7, 11, 8, 11, 17, 17, 10, 15, 20, 8, 14, 19, 18]\n",
      "pop -> 4 [4, 6, 6, 7, 8, 16, 8, 7, 11, 18, 11, 17, 17, 10, 15, 20, 8, 14, 19]\n",
      "pop -> 4 [6, 6, 8, 7, 8, 16, 10, 7, 11, 18, 11, 17, 17, 19, 15, 20, 8, 14]\n",
      "pop -> 6 [6, 7, 8, 7, 8, 16, 10, 8, 11, 18, 11, 17, 17, 19, 15, 20, 14]\n",
      "pop -> 6 [7, 7, 8, 8, 8, 16, 10, 14, 11, 18, 11, 17, 17, 19, 15, 20]\n",
      "pop -> 7 [7, 8, 8, 8, 11, 16, 10, 14, 11, 18, 20, 17, 17, 19, 15]\n",
      "pop -> 7 [8, 8, 10, 8, 11, 16, 15, 14, 11, 18, 20, 17, 17, 19]\n",
      "pop -> 8 [8, 8, 10, 11, 11, 16, 15, 14, 19, 18, 20, 17, 17]\n",
      "pop -> 8 [8, 11, 10, 11, 17, 16, 15, 14, 19, 18, 20, 17]\n",
      "pop -> 8 [10, 11, 15, 11, 17, 16, 17, 14, 19, 18, 20]\n",
      "pop -> 10 [11, 11, 15, 14, 17, 16, 17, 20, 19, 18]\n",
      "pop -> 11 [11, 14, 15, 18, 17, 16, 17, 20, 19]\n",
      "pop -> 11 [14, 17, 15, 18, 19, 16, 17, 20]\n",
      "pop -> 14 [15, 17, 16, 18, 19, 20, 17]\n",
      "pop -> 15 [16, 17, 17, 18, 19, 20]\n",
      "pop -> 16 [17, 17, 20, 18, 19]\n",
      "pop -> 17 [17, 18, 20, 19]\n",
      "pop -> 17 [18, 19, 20]\n",
      "pop -> 18 [19, 20]\n",
      "pop -> 19 [20]\n",
      "pop -> 20 []\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "heap = []\n",
    "for _ in range(25):\n",
    "    element = random.randint(1, 20)\n",
    "    push(heap, element)\n",
    "    print(\"push\", element, \"->\", heap)\n",
    "\n",
    "while heap:\n",
    "    print(\"pop ->\", pop(heap), heap)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}