cp_library/alg/iter/cmpr/icoord_compress_with_queries_fn.py

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

Depends on

Verified with

test/library-checker/data-structure/static_rectangle_add_rectangle_sum_bit_monoid.test.py

Code

import cp_library.__header__
import cp_library.alg.__header__
from cp_library.alg.dp.max2_fn import max2
import cp_library.alg.iter.__header__
import cp_library.alg.iter.cmpr.__header__

def icoord_compress_with_queries(*A: list[int], x=0, distinct=False):
    N = mx = 0
    for Ai in A: N += len(Ai); mx = max2(mx, len(Ai))
    si, mi = pack_sm(mx-1); sj, mj = pack_sm((len(A)-1)<<si)
    S, k = [0]*N, 0
    for i,Ai in enumerate(A):
        for j,a in enumerate(Ai): S[k]=a << sj | i << si | j; k += 1
    S.sort(); r = p = -1
    for aji in S:
        a, i, j = aji >> sj, (aji&mj) >> si , aji & mi
        if x<=i and (distinct or a != p): r = r+1; p = a
        A[i][j] = r+(i<x)
    return A
from cp_library.bit.pack.pack_sm_fn import pack_sm

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



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



def icoord_compress_with_queries(*A: list[int], x=0, distinct=False):
    N = mx = 0
    for Ai in A: N += len(Ai); mx = max2(mx, len(Ai))
    si, mi = pack_sm(mx-1); sj, mj = pack_sm((len(A)-1)<<si)
    S, k = [0]*N, 0
    for i,Ai in enumerate(A):
        for j,a in enumerate(Ai): S[k]=a << sj | i << si | j; k += 1
    S.sort(); r = p = -1
    for aji in S:
        a, i, j = aji >> sj, (aji&mj) >> si , aji & mi
        if x<=i and (distinct or a != p): r = r+1; p = a
        A[i][j] = r+(i<x)
    return A


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