cp_library/alg/divcon/qselect_fn.py

• View this file on GitHub
• Last update: 2025-07-11 23:11:42+09:00

Depends on

• cp_library/alg/divcon/median_of_three_fn.py
• cp_library/alg/divcon/partition_fn.py

Required by

• cp_library/math/median_fn.py

Verified with

• test/atcoder/arc/arc122_b_insurance_median.test.py
• test/atcoder/arc/arc182_d_increment_decrement_again_qselect.test.py

Code

import cp_library.alg.divcon.__header__
from cp_library.alg.divcon.median_of_three_fn import median_of_three
from cp_library.alg.divcon.partition_fn import partition

def qselect(A, k, l=0, r=None):
    '''Find kth element in subarray [l,r)'''
    if r is None: r = len(A)
    while l != r-1:
        if k < (p := partition(A, l, r, median_of_three(A,l,r))): r = p
        elif k > p: l = p+1
        else: return A[k]
    return A[k]

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''

def median_of_three(A, l, r):
    '''Select pivot as median of first, middle, and last elements'''
    if r - l < 3: return l
    mid = (l+r) >> 1
    if A[mid] < A[l]:
        A[l], A[mid] = A[mid], A[l]
    if A[r-1] < A[mid]:
        A[mid], A[r-1] = A[r-1], A[mid]
        if A[mid] < A[l]:
            A[l], A[mid] = A[mid], A[l]
    return mid

def partition(A, l, r, p) -> int:
    '''Partition subarray [l,r)'''
    A[p], A[r], p = A[r := r-1], A[p], l
    for j in range(l, r):
        if A[j] <= A[r]: A[p], A[j], p = A[j], A[p], p+1
    A[p], A[r] = A[r], A[p]
    return p

def qselect(A, k, l=0, r=None):
    '''Find kth element in subarray [l,r)'''
    if r is None: r = len(A)
    while l != r-1:
        if k < (p := partition(A, l, r, median_of_three(A,l,r))): r = p
        elif k > p: l = p+1
        else: return A[k]
    return A[k]

Back to top page