cp_library/alg/iter/sort_parallel_fn.py

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

Depends on

Code

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

def sort_parallel(*L: list, reverse=False):
    inv, order = [-1]*(N := len(L[0])), argsort(L[0], reverse=reverse)
    for i, idx in enumerate(order): inv[idx] = i
    for j in range(N):
        i, k = inv[j], order[j]
        for A in L: A[j], A[k] = A[k], A[j] 
        order[i], inv[k] = k, i
    return L

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


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

def pack_enc(a: int, b: int, s: int):
    return a << s | b
    
def pack_dec(ab: int, s: int, m: int):
    return ab >> s, ab & m

def pack_indices(A, s):
    return [a << s | i for i,a in enumerate(A)]

def argsort(A: list[int], reverse=False):
    s, m = pack_sm(len(A))
    if reverse:
        I = [a<<s|i^m for i,a in enumerate(A)]
        I.sort(reverse=True)
        for i,ai in enumerate(I): I[i] = (ai^m)&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 sort_parallel(*L: list, reverse=False):
    inv, order = [-1]*(N := len(L[0])), argsort(L[0], reverse=reverse)
    for i, idx in enumerate(order): inv[idx] = i
    for j in range(N):
        i, k = inv[j], order[j]
        for A in L: A[j], A[k] = A[k], A[j] 
        order[i], inv[k] = k, i
    return L