{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c84215e04e1c",
   "metadata": {},
   "source": [
    "# Top K Frequent Elements\n",
    "\n",
    "| Field      | Value                                   |\n",
    "|------------|-----------------------------------------|\n",
    "| Source     | [link](https://leetcode.com/problems/top-k-frequent-elements/)  |\n",
    "| Difficulty | Medium                    |\n",
    "| Topics     | `array`, `hash map`, `bucket sort`, `heap`                     |\n",
    "| Date       | 2026-10-08                              |\n",
    "| Status     | Solved                                  |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f59b8cccd29d",
   "metadata": {},
   "source": [
    "## 1. Problem\n",
    "\n",
    "> Given an integer array `nums` and an integer `k`, return the `k` most frequent elements within\n",
    "> the array. The test cases are generated such that the answer is always unique.\n",
    "> You may return the output in any order.\n",
    "\n",
    "**Example 1**\n",
    "```\n",
    "Input:  nums = [1, 2, 2, 3, 3, 3], k = 2\n",
    "Output: [2, 3]\n",
    "```\n",
    "\n",
    "**Example 2**\n",
    "```\n",
    "Input:  nums = [7, 7], k = 1\n",
    "Output: [7]\n",
    "```\n",
    "\n",
    "**Constraints**\n",
    "- `1 <= len(nums) <= 10^4`\n",
    "- `-1000 <= nums[i] <= 1000`\n",
    "- `1 <= k <=` number of distinct elements in `nums`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbac771a374a",
   "metadata": {},
   "source": [
    "## 2. Thinking\n",
    "\n",
    "Let $n$ = `len(nums)`, $u$ = number of distinct values ($u \\le n$, and here $u \\le 2001$).\n",
    "\n",
    "- **Brute force idea:** count every value, then sort the $u$ distinct values by count and take the\n",
    "  first $k$: $O(n + u \\log u)$. That's already fine; the follow-up asks for better than $O(n \\log n)$.\n",
    "- **Key insight:** step 1 is always **count with a hash map** ($O(n)$). Step 2 is \"pick the top $k$\n",
    "  by count\", and there are three ways to do it:\n",
    "  - **heap** of size $k$: $O(u \\log k)$;\n",
    "  - **bucket sort**: a count can only be between 1 and $n$, so put each value in `buckets[count]`\n",
    "    and walk from the highest count down: $O(n)$;\n",
    "  - full sort: $O(u \\log u)$.\n",
    "- **Unique answer guarantee:** the bucket that takes you past $k$ can never have a tie, so you can\n",
    "  add whole buckets with `extend` and stop when `len(result) >= k`.\n",
    "- **Edge cases to remember:** one element (`[7]`), every element the same (count $= n$, so the\n",
    "  bucket list needs index $n$), $k$ = number of distinct values, negative numbers."
   ]
  },
  {
   "cell_type": "code",
   "id": "939f330661fd",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Setup: run this cell first. Lets you run pytest tests inside the notebook.\n",
    "try:\n",
    "    import ipytest\n",
    "except ImportError:  # Colab and other fresh kernels don't ship ipytest\n",
    "    %pip install -q ipytest\n",
    "    import ipytest\n",
    "import pytest\n",
    "\n",
    "ipytest.autoconfig()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f5d37b2b8fe",
   "metadata": {},
   "source": [
    "## 3. Solution\n",
    "\n",
    "### Approach A — my version (bucket sort, after fixing the two bugs)\n",
    "\n",
    "**Complexity:** Time $O(n)$ · Space $O(n)$. The Big-O is already optimal. It's slow in\n",
    "practice because of constant factors:\n",
    "1. **`n + 1` buckets, sized by `len(nums)`.** With $n = 10^4$ that is 10,001 Python list objects\n",
    "   (~700 KiB), built in a Python loop, even though the highest count is often around 5. This is\n",
    "   the main memory cost and about a quarter of the runtime.\n",
    "2. **The bucket walk never stops.** After `result` has $k$ elements the `if` keeps failing, but\n",
    "   the loop still visits all 10,001 buckets. `temp[::-1]` also copies the whole list first.\n",
    "3. **Counting in a Python loop.** `frequency[num] += 1` runs $n$ times in bytecode;\n",
    "   `Counter(nums)` does the same counting in C, about 2× faster.\n",
    "4. `int(x)` is left over from the old `str` keys and does nothing now.\n",
    "\n",
    "The code is unchanged except for the signature and comments."
   ]
  },
  {
   "cell_type": "code",
   "id": "d8fe2907be2f",
   "metadata": {},
   "source": [
    "from collections import Counter, defaultdict\n",
    "import heapq\n",
    "\n",
    "def top_k_frequent_mine(nums: list[int], k: int) -> list[int]:\n",
    "    frequency = defaultdict(int)\n",
    "    temp = [[] for _ in range(len(nums) + 1)]  # SLOW: n+1 lists, most of them stay empty\n",
    "    for num in nums:\n",
    "        frequency[num] += 1                      # SLOW: counting loop in Python, Counter is C\n",
    "    for x, y in frequency.items():\n",
    "        temp[y].append(int(x))                   # int() does nothing now\n",
    "    result = []\n",
    "    for z in temp[::-1]:                         # copies the list, then visits every bucket\n",
    "        if len(result) < k:\n",
    "            result.extend(z)                     # no break: keeps looping after it has k\n",
    "    return result"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "b71b417f444c",
   "metadata": {},
   "source": [
    "### Approach B — bucket sort, done efficiently (the $O(n)$ interview answer)\n",
    "\n",
    "**Complexity:** Time $O(n)$ · Space $O(n)$ worst case, but only $O(u + \\text{max count})$ in practice.\n",
    "Same idea as mine with three fixes: count with `Counter`, make only `max count + 1` buckets, and\n",
    "`break` once there are $k$ results. This alone makes it ~3× faster and cuts peak memory ~7×."
   ]
  },
  {
   "cell_type": "code",
   "id": "06a1d908f26e",
   "metadata": {},
   "source": [
    "def top_k_frequent_bucket(nums: list[int], k: int) -> list[int]:\n",
    "    count = Counter(nums)\n",
    "    buckets = [[] for _ in range(max(count.values()) + 1)]  # index = count\n",
    "    for num, c in count.items():\n",
    "        buckets[c].append(num)\n",
    "    result = []\n",
    "    for bucket in reversed(buckets):  # highest count first, no copy\n",
    "        result.extend(bucket)\n",
    "        if len(result) >= k:\n",
    "            break\n",
    "    return result"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "0afc22113885",
   "metadata": {},
   "source": [
    "### Approach C — heap of size $k$ (`heapq.nlargest`)\n",
    "\n",
    "**Complexity:** Time $O(n + u \\log k)$ · Space $O(u + k)$.\n",
    "Count, then keep only the $k$ largest counts in a min-heap. `nlargest` does this in C. It's the\n",
    "answer to give when the data is a stream or $k$ is much smaller than $u$."
   ]
  },
  {
   "cell_type": "code",
   "id": "f0d54afaad4c",
   "metadata": {},
   "source": [
    "def top_k_frequent_heap(nums: list[int], k: int) -> list[int]:\n",
    "    count = Counter(nums)\n",
    "    return heapq.nlargest(k, count, key=count.get)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "725a71b05470",
   "metadata": {},
   "source": [
    "### Approach D — `Counter.most_common` ✅ best practice (Pythonic)\n",
    "\n",
    "**Complexity:** Time $O(n + u \\log k)$ · Space $O(u + k)$.\n",
    "`most_common(k)` calls `heapq.nlargest` internally, so this is Approach C in one line. In\n",
    "everyday Python, write this. In an interview, write it, then explain B (the $O(n)$ bucket sort)\n",
    "when asked for better than $O(n \\log n)$."
   ]
  },
  {
   "cell_type": "code",
   "id": "190fea063b32",
   "metadata": {},
   "source": [
    "def top_k_frequent(nums: list[int], k: int) -> list[int]:\n",
    "    return [num for num, _ in Counter(nums).most_common(k)]"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "7425c5edb1d4",
   "metadata": {},
   "source": [
    "## 4. Tests\n",
    "\n",
    "Any output order is allowed, so results are compared with `sorted`. Every approach runs against\n",
    "every case."
   ]
  },
  {
   "cell_type": "code",
   "id": "6ff2659119d4",
   "metadata": {},
   "source": [
    "%%ipytest -q\n",
    "\n",
    "APPROACHES = [top_k_frequent_mine, top_k_frequent_bucket, top_k_frequent_heap, top_k_frequent]\n",
    "\n",
    "CASES = [\n",
    "    ([1, 2, 2, 3, 3, 3], 2, [2, 3]),\n",
    "    ([7, 7], 1, [7]),                          # count == len(nums): needs bucket index n\n",
    "    ([7], 1, [7]),                             # single element\n",
    "    ([4, 4, 4, 4], 1, [4]),                    # all the same\n",
    "    ([1, 2], 2, [1, 2]),                       # k == number of distinct values, all tied\n",
    "    ([-1, -1, -1000, 1000, 1000, 1000], 2, [-1, 1000]),  # negative numbers\n",
    "    ([5, 1, 5, 2, 5, 2, 3], 2, [5, 2]),        # answer not in input order\n",
    "    ([1] * 10_000, 1, [1]),                    # max size, one value\n",
    "]\n",
    "\n",
    "@pytest.mark.parametrize(\"fn\", APPROACHES)\n",
    "@pytest.mark.parametrize(\"nums, k, expected\", CASES)\n",
    "def test_top_k_frequent(fn, nums, k, expected):\n",
    "    assert sorted(fn(nums, k)) == sorted(expected)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "aeb8fb7cb9fd",
   "metadata": {},
   "source": [
    "## 5. Scratch / Experiments\n",
    "\n",
    "Worst-ish case at the max size: $10^4$ random values from the full range $[-1000, 1000]$,\n",
    "so $u \\approx 2001$ and the highest count is only ~15.\n",
    "\n",
    "All four are $O(n)$ or close to it, so the gaps are **constant factors**:\n",
    "- **A vs B:** same algorithm. B is ~3× faster only because it counts in C, builds ~15 buckets\n",
    "  instead of 10,001, and stops early.\n",
    "- **B vs C/D:** bucket sort is $O(n)$ and the heap is $O(n + u \\log k)$, but with $k$ small,\n",
    "  $\\log k$ is tiny and `nlargest` runs in C, so they finish about even. Big-O tells you how it\n",
    "  grows, not which is faster at $n = 10^4$.\n",
    "- The second block measures peak memory: A's 10,001 empty lists are almost all of it.\n",
    "\n",
    "LeetCode's runtime/memory percentiles are noisy (±20% between runs of the same code), so treat\n",
    "them as a hint, not a verdict."
   ]
  },
  {
   "cell_type": "code",
   "id": "0b7ed0affaaa",
   "metadata": {},
   "source": [
    "import random, tracemalloc\n",
    "\n",
    "random.seed(0)\n",
    "nums = [random.randint(-1000, 1000) for _ in range(10_000)]\n",
    "k = 10\n",
    "\n",
    "for fn in APPROACHES:\n",
    "    print(fn.__name__)\n",
    "    %timeit fn(nums, k)\n",
    "\n",
    "for fn in APPROACHES:\n",
    "    tracemalloc.start()\n",
    "    fn(nums, k)\n",
    "    peak = tracemalloc.get_traced_memory()[1]\n",
    "    tracemalloc.stop()\n",
    "    print(f\"{fn.__name__:24s} peak {peak / 1024:6.0f} KiB\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "bec8bed23a56",
   "metadata": {},
   "source": [
    "## 6. Takeaways\n",
    "\n",
    "- **Pattern:** \"top $k$ / $k$ most frequent / $k$ largest\" → **count with a hash map, then\n",
    "  select**: `Counter.most_common(k)` / `heapq.nlargest` ($O(u \\log k)$), or **bucket sort by\n",
    "  count** when the counts are bounded by $n$ ($O(n)$).\n",
    "- **Mistakes I made:**\n",
    "  - Sized the bucket list as `len(nums)`: a value can appear `len(nums)` times, so I needed\n",
    "    `len(nums) + 1` (crashed on `[7, 7]`).\n",
    "  - Used `<= k` as \"still need more\": that's `< k`. Off-by-one in a stop condition: trace the\n",
    "    example by hand and check the last iteration.\n",
    "  - Used `str(num)` as the dict key: ints are hashable, so convert nothing.\n",
    "  - **Optimal Big-O can still be slow.** Mine was $O(n)$ but allocated 10,001 lists and never\n",
    "    stopped looping. Size buffers by what's needed (`max(count.values())`), `break` when done,\n",
    "    and prefer C-implemented builtins (`Counter`, `heapq`) over hand-written loops.\n",
    "- **Related problems:** Kth Largest Element in an Array (#215, heap / quickselect for top $k$),\n",
    "  Top K Frequent Words (#692, same idea plus a tie-break by word), Sort Characters By Frequency\n",
    "  (#451, count + bucket sort), K Closest Points to Origin (#973, `heapq.nsmallest` with a key),\n",
    "  Group Anagrams (#49, the counting-with-a-dict step)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}