Find the n smallest elements in a dataset. Equivalent to: sorted(iterable, key=key)[:n]
(n, iterable, key=None)
| 475 | # http://code.activestate.com/recipes/577573-compare-algorithms-for-heapqsmallest |
| 476 | |
| 477 | def nsmallest(n, iterable, key=None): |
| 478 | """Find the n smallest elements in a dataset. |
| 479 | |
| 480 | Equivalent to: sorted(iterable, key=key)[:n] |
| 481 | """ |
| 482 | |
| 483 | # Short-cut for n==1 is to use min() |
| 484 | if n == 1: |
| 485 | it = iter(iterable) |
| 486 | sentinel = object() |
| 487 | result = min(it, default=sentinel, key=key) |
| 488 | return [] if result is sentinel else [result] |
| 489 | |
| 490 | # When n>=size, it's faster to use sorted() |
| 491 | try: |
| 492 | size = len(iterable) |
| 493 | except (TypeError, AttributeError): |
| 494 | pass |
| 495 | else: |
| 496 | if n >= size: |
| 497 | return sorted(iterable, key=key)[:n] |
| 498 | |
| 499 | # When key is none, use simpler decoration |
| 500 | if key is None: |
| 501 | it = iter(iterable) |
| 502 | # put the range(n) first so that zip() doesn't |
| 503 | # consume one too many elements from the iterator |
| 504 | result = [(elem, i) for i, elem in zip(range(n), it)] |
| 505 | if not result: |
| 506 | return result |
| 507 | heapify_max(result) |
| 508 | top = result[0][0] |
| 509 | order = n |
| 510 | _heapreplace = heapreplace_max |
| 511 | for elem in it: |
| 512 | if elem < top: |
| 513 | _heapreplace(result, (elem, order)) |
| 514 | top, _order = result[0] |
| 515 | order += 1 |
| 516 | result.sort() |
| 517 | return [elem for (elem, order) in result] |
| 518 | |
| 519 | # General case, slowest method |
| 520 | it = iter(iterable) |
| 521 | result = [(key(elem), i, elem) for i, elem in zip(range(n), it)] |
| 522 | if not result: |
| 523 | return result |
| 524 | heapify_max(result) |
| 525 | top = result[0][0] |
| 526 | order = n |
| 527 | _heapreplace = heapreplace_max |
| 528 | for elem in it: |
| 529 | k = key(elem) |
| 530 | if k < top: |
| 531 | _heapreplace(result, (k, order, elem)) |
| 532 | top, _order, _elem = result[0] |
| 533 | order += 1 |
| 534 | result.sort() |