mirror of
https://github.com/TheAlgorithms/Python.git
synced 2024-12-05 02:40:16 +00:00
bc8df6de31
* [pre-commit.ci] pre-commit autoupdate updates: - [github.com/astral-sh/ruff-pre-commit: v0.2.2 → v0.3.2](https://github.com/astral-sh/ruff-pre-commit/compare/v0.2.2...v0.3.2) - [github.com/pre-commit/mirrors-mypy: v1.8.0 → v1.9.0](https://github.com/pre-commit/mirrors-mypy/compare/v1.8.0...v1.9.0) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
"""
|
|
Find the kth smallest element in linear time using divide and conquer.
|
|
Recall we can do this trivially in O(nlogn) time. Sort the list and
|
|
access kth element in constant time.
|
|
|
|
This is a divide and conquer algorithm that can find a solution in O(n) time.
|
|
|
|
For more information of this algorithm:
|
|
https://web.stanford.edu/class/archive/cs/cs161/cs161.1138/lectures/08/Small08.pdf
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from random import choice
|
|
|
|
|
|
def random_pivot(lst):
|
|
"""
|
|
Choose a random pivot for the list.
|
|
We can use a more sophisticated algorithm here, such as the median-of-medians
|
|
algorithm.
|
|
"""
|
|
return choice(lst)
|
|
|
|
|
|
def kth_number(lst: list[int], k: int) -> int:
|
|
"""
|
|
Return the kth smallest number in lst.
|
|
>>> kth_number([2, 1, 3, 4, 5], 3)
|
|
3
|
|
>>> kth_number([2, 1, 3, 4, 5], 1)
|
|
1
|
|
>>> kth_number([2, 1, 3, 4, 5], 5)
|
|
5
|
|
>>> kth_number([3, 2, 5, 6, 7, 8], 2)
|
|
3
|
|
>>> kth_number([25, 21, 98, 100, 76, 22, 43, 60, 89, 87], 4)
|
|
43
|
|
"""
|
|
# pick a pivot and separate into list based on pivot.
|
|
pivot = random_pivot(lst)
|
|
|
|
# partition based on pivot
|
|
# linear time
|
|
small = [e for e in lst if e < pivot]
|
|
big = [e for e in lst if e > pivot]
|
|
|
|
# if we get lucky, pivot might be the element we want.
|
|
# we can easily see this:
|
|
# small (elements smaller than k)
|
|
# + pivot (kth element)
|
|
# + big (elements larger than k)
|
|
if len(small) == k - 1:
|
|
return pivot
|
|
# pivot is in elements bigger than k
|
|
elif len(small) < k - 1:
|
|
return kth_number(big, k - len(small) - 1)
|
|
# pivot is in elements smaller than k
|
|
else:
|
|
return kth_number(small, k)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import doctest
|
|
|
|
doctest.testmod()
|