O(N log N)复杂度py正经题解

先初始化一个大小为 $n + 2$ 的数组 tree
然后更新树状数组中的值
每次更新时，idx 会加上 idx & -idx，直到 idx 超过 size
最坏情况下，idx 要遍历到 size
然后查询前缀和
每次查询时，idx 会减去 idx & -idx，直到 idx 小于等于 0
最坏情况下，idx 要遍历到 1
找第 k 个元素的位置，随后用二分查找的方法，通过位掩码来定位位置
然后用正则表达式提取字符串中的所有整数
将数组格式化为特定字符串格式
读取输入并分割成数据列表
初始化 B 数组和 bit 对象
遍历 arr 进行更新和查询操作
初始化 A 数组和 bit 对象
遍历 arr 进行更新和查找操作
对于类型 'A' 和 'B'，主操作都基于 BIT 类的 update, query, 和 find_kth，这些方法的时间复杂度都是 $O(log N)$
所以整体时间复杂度为$O(N log N)$

题解：

import sys 
import re  
class BIT: 
    def __init__(self, n): 
        self.size  = n 
        self.tree  = [0] * (n + 2) 
    def update(self, idx, delta=1): 
        while idx <= self.size:  
            self.tree[idx]  += delta 
            idx += idx & -idx 
    def query(self, idx): 
        res = 0 
        while idx > 0: 
            res += self.tree[idx]  
            idx -= idx & -idx 
        return res 
    def find_kth(self, k): 
        idx = 0 
        bitMask = 1 
        while bitMask <= self.size:  
            bitMask <<= 1 
        bitMask >>= 1 
        while bitMask > 0: 
            temp = idx + bitMask 
            if temp <= self.size  and self.tree[temp]  < k: 
                idx = temp 
                k -= self.tree[temp]  
            bitMask >>= 1 
        return idx + 1  
def parse_array(s): 
    arr = [] 
    nums = re.findall(r'-?\d+',  s) 
    for num in nums: 
        arr.append(int(num))  
    return arr  
def format_output(type, arr): 
    res = f"{type}=(" 
    for i, val in enumerate(arr): 
        res += str(val) 
        if i != len(arr) - 1: 
            res += "," 
    res += ")" 
    return res  
def main(): 
    input = sys.stdin.read  
    data = input().split() 
    N = int(data[0]) 
    line = data[1] 
    type = line[0] 
    arr = parse_array(line)  
    if type == 'A': 
        B = [0] * N 
        bit = BIT(N) 
        B[0] = 0 
        bit.update(arr[0]  + 1) 
        for i in range(1, N): 
            if arr[i] == 0: 
                B[i] = 0 
            else: 
                B[i] = bit.query(arr[i])  
            bit.update(arr[i]  + 1) 
        print(format_output('B', B)) 
    else: 
        A = [0] * N 
        bit = BIT(N) 
        for i in range(N): 
            bit.update(i  + 1) 
        for i in range(N - 1, -1, -1): 
            k = arr[i] + 1 
            pos = bit.find_kth(k)  
            A[i] = pos - 1 
            bit.update(pos,  -1) 
        print(format_output('A', A))  
if __name__ == "__main__": 
    main()