cp_library/alg/divcon/qselect_fn.py

• View this file on GitHub
• Last update: 2025-03-30 20:17:47+09:00

Depends on

• 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 random import randint
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, randint(l,r-1))): r = p
        elif k > p: l = p+1
        else: return A[k]
    return A[k]

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''
from random import randint

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, randint(l,r-1))): r = p
        elif k > p: l = p+1
        else: return A[k]
    return A[k]

Back to top page