cp_library/alg/iter/rank/rank_multi_fn.py

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

Depends on

Code

import cp_library.__header__
import cp_library.alg.__header__
import cp_library.alg.iter.__header__
import cp_library.alg.iter.rank.__header__
from cp_library.alg.iter.rank.irank_multi_fn import irank

def rank(*A: list[int], distinct = False): return *(R := tuple(Ai.copy() for Ai in A)), irank(*R, distinct)

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





def max2(a, b):
    return a if a > b else b

def irank(*A: list[int], distinct = False):
    N = mxj = 0
    for Ai in A: N += len(Ai); mxj = max2(mxj, len(Ai))
    P = Packer3(len(A)-1, mxj); V = P.enumerate(A, N); V.sort()
    if distinct:
        for r,aij in enumerate(V):a,i,j=P.dec(aij);A[i][j],V[r]=r,a
    elif V:
        r, p = -1, V[-1]+1 # set p to unique value to trigger `if a != p` on first elm
        for aij in V:
            a,i,j=P.dec(aij)
            if a!=p:V[r:=r+1]=p=a
            A[i][j]=r
        del V[r+1:]
    return V



class Packer3:
    def __init__(P, mxb: int, mxc: int):
        bb, bc = mxb.bit_length(), mxc.bit_length()
        P.mc, P.mb, P.sb, P.sa = (1<<bc)-1, (1<<bb)-1, bc, bc+bb
    def enc(P, a: int, b: int, c: int): return a << P.sa | b << P.sb | c
    def dec(P, x: int) -> tuple[int, int, int]: return x >> P.sa, (x >> P.sb) & P.mb, x & P.mc
    def enumerate(P, A, N, reverse=False): 
        V, k = [0]*N, 0
        if reverse:
            for i,Ai in enumerate(A):
                for j, a in enumerate(Ai):V[k]=P.enc(-a, i, j);k+=1
        else:
            for i,Ai in enumerate(A):
                for j, a in enumerate(Ai):V[k]=P.enc(a, i, j);k+=1
        return V

def rank(*A: list[int], distinct = False): return *(R := tuple(Ai.copy() for Ai in A)), irank(*R, distinct)