cp_library/alg/iter/arg/argsort_bounded_fn.py

• View this file on GitHub
• Last update: 2025-06-20 03:24:59+09:00

Depends on

• argsort (cp_library/alg/iter/arg/argsort_fn.py)
• cp_library/bit/pack/pack_sm_fn.py

Verified with

• test/library-checker/graph/incremental_scc.test.py

Code

import cp_library.__header__
import cp_library.alg.__header__
import cp_library.alg.iter.__header__
import cp_library.alg.iter.arg.__header__
from cp_library.alg.iter.arg.argsort_fn import argsort

def argsort_bounded(A, mx=None, reverse=False):
    N = len(A)
    if mx is None: mx = max(A)
    if N*N.bit_length() < mx or mx < 1000: return argsort(A, reverse)
    I, cnt, t = [0]*N, [0]*(mx+1), 0
    for a in A: cnt[a] += 1
    if reverse:
        for a in range(mx+1): cnt[~a], t = t, t+cnt[~a]
    else:
        for a in range(mx+1): cnt[a], t = t, t+cnt[a]
    for i,a in enumerate(A): I[cnt[a]] = i; cnt[a] += 1
    return I

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




def argsort(A: list[int], reverse=False):
    s, m = pack_sm(len(A))
    if reverse:
        I = [a<<s|m^i for i,a in enumerate(A)]
        I.sort(reverse=True)
        for i,ai in enumerate(I): I[i] = m^ai&m
    else:
        I = [a<<s|i for i,a in enumerate(A)]
        I.sort()
        for i,ai in enumerate(I): I[i] = ai&m
    return I


def pack_sm(N: int): s=N.bit_length(); return s,(1<<s)-1

def argsort_bounded(A, mx=None, reverse=False):
    N = len(A)
    if mx is None: mx = max(A)
    if N*N.bit_length() < mx or mx < 1000: return argsort(A, reverse)
    I, cnt, t = [0]*N, [0]*(mx+1), 0
    for a in A: cnt[a] += 1
    if reverse:
        for a in range(mx+1): cnt[~a], t = t, t+cnt[~a]
    else:
        for a in range(mx+1): cnt[a], t = t, t+cnt[a]
    for i,a in enumerate(A): I[cnt[a]] = i; cnt[a] += 1
    return I

Back to top page