#561220 (PyPy3) No.1248 Kth Sum

提出ソース

結果

問題	No.1248 Kth Sum
ユーザー	chineristAC
提出日時	2020-10-02 22:01:32
言語	PyPy3 (7.3.15)
結果	WA
実行時間	-
コード長	8,572 bytes
コンパイル時間	245 ms
コンパイル使用メモリ	82,568 KB
実行使用メモリ	132,864 KB
最終ジャッジ日時	2024-07-17 13:20:01
合計ジャッジ時間	9,718 ms
ジャッジサーバーID （参考情報）	judge5 / judge4

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	38 ms 53,948 KB
testcase_01	AC	38 ms 54,032 KB
testcase_02	AC	39 ms 53,852 KB
testcase_03	AC	38 ms 54,168 KB
testcase_04	AC	38 ms 53,544 KB
testcase_05	WA	-
testcase_06	WA	-
testcase_07	AC	38 ms 54,352 KB
testcase_08	WA	-
testcase_09	WA	-
testcase_10	WA	-
testcase_11	WA	-
testcase_12	WA	-
testcase_13	AC	92 ms 99,644 KB
testcase_14	WA	-
testcase_15	AC	93 ms 99,396 KB
testcase_16	WA	-
testcase_17	WA	-
testcase_18	WA	-
testcase_19	WA	-
testcase_20	WA	-
testcase_21	WA	-
testcase_22	WA	-
testcase_23	AC	37 ms 54,840 KB
testcase_24	WA	-
testcase_25	AC	275 ms 109,008 KB
testcase_26	WA	-
testcase_27	AC	94 ms 99,440 KB
testcase_28	AC	94 ms 99,468 KB
testcase_29	AC	93 ms 99,568 KB
testcase_30	WA	-
testcase_31	WA	-
testcase_32	WA	-
testcase_33	WA	-
testcase_34	WA	-
testcase_35	WA	-
testcase_36	WA	-
testcase_37	WA	-

権限があれば一括ダウンロードができます

ソースコード

raw source code

"""
多重平衡二分木
挿入、削除、bisect、k番目の値の取得,k番目までの値の和の取得がO(logL)で可能
同じ値には違うidを与えて区別する
idが足りなくなったらid数を二倍にして保持している値を一から挿入しなおすことでならし計算量がO(logL+logK)になるはず…
あらかじめid数を設定することもできる

verify:ABC127F,ABC170E
"""
class MultiBalancingTree:
    def __init__(self, n,id=0):
        self.N = n
        self.id_size=id
        self.root = self.node(1<<(n+self.id_size), 1<<(n+self.id_size),self.id_size)
        self.val_dic={}

    def debug(self):
        def debug_info(nd_):
            return (nd_.value,(nd_.value - 1)>>self.id_size, nd_.pivot , nd_.left.value - 1 if nd_.left else -1, nd_.right.value - 1 if nd_.right else -1,nd_.cnt)

        def debug_node(nd):
            re = []
            if nd.left:
                re += debug_node(nd.left)
            if nd.value: re.append(debug_info(nd))
            if nd.right:
                re += debug_node(nd.right)
            return re
        print("Debug - root =", self.root.value - 1, debug_node(self.root)[:50])

    def rebuild(self,v):
        self.id_size+=1
        #print("rebuild:",self.id_size)
        self.root=self.node(1<<(self.N+self.id_size), 1<<(self.N+self.id_size),self.id_size)
        check=[]
        for val in self.val_dic:
            cnt=self.val_dic[val]
            self.val_dic[val]=0
            for j in range(cnt):
                self.append(val)
        self.append(v)

    def append(self,v):
        if v not in self.val_dic:
            self.val_dic[v]=0
        id=self.val_dic[v]
        if id>=2**self.id_size:
            self.rebuild(v)
        else:
            V=(v<<self.id_size)+id
            self.val_dic[v]+=1
            self._append(V)

    def _append(self, v):
        v += 1
        nd = self.root
        while True:
            if v == nd.value:
                return 0
            else:
                nd.cnt+=1
                nd.sum+=(v-1)>>self.id_size
                mi, ma = min(v, nd.value), max(v, nd.value)
                if mi < nd.pivot:
                    nd.value = ma
                    if nd.left:
                        nd = nd.left
                        v = mi
                    else:
                        p = nd.pivot
                        nd.left = self.node(mi, p - (p&-p)//2,self.id_size)
                        break
                else:
                    nd.value = mi
                    if nd.right:
                        nd = nd.right
                        v = ma
                    else:
                        p = nd.pivot
                        nd.right = self.node(ma, p + (p&-p)//2,self.id_size)
                        break

    def kth(self,k):
        cnt=k
        nd=self.root
        while True:
            if nd.left:
                if nd.left.cnt<k:
                    k-=nd.left.cnt+1
                    if nd.right:
                        nd=nd.right
                    else:
                        raise IndexError("BalancingTree doesn't hold kth number")
                elif nd.left.cnt==k:
                    return (nd.value-1)>>self.id_size
                else:
                    nd=nd.left
            else:
                if k==0:
                    return (nd.value-1)>>self.id_size
                if nd.right:
                    k-=1
                    nd=nd.right
                else:
                    raise IndexError("BalancingTree doesn't hold kth number")

    def kth_sum(self,k):
        if k==-1:
            return 0
        cnt=k
        res_sum=0
        nd=self.root
        while True:
            if nd.left:
                if nd.left.cnt<k:
                    k-=nd.left.cnt+1
                    res_sum+=nd.left.sum+((nd.value-1)>>self.id_size)
                    if nd.right:
                        nd=nd.right
                    else:
                        raise IndexError("BalancingTree doesn't hold kth number")
                elif nd.left.cnt==k:
                    res_sum+=nd.left.sum+((nd.value-1)>>self.id_size)
                    return res_sum,(nd.value-1)>>self.id_size
                else:
                    nd=nd.left
            else:
                if k==0:
                    return res_sum+((nd.value-1)>>self.id_size),(nd.value-1)>>self.id_size
                if nd.right:
                    k-=1
                    res_sum+=((nd.value-1)>>self.id_size)
                    nd=nd.right
                else:
                    raise IndexError("BalancingTree doesn't hold kth number")

    def kth_sum_big(self,k):
        res_sum=self.sum
        if self.size>k:
            res_sum-=self.kth_sum(self.size-k-2)
            return res_sum
        else:
            raise IndexError("BalancingTree doesn't hold kth number")

    def leftmost(self, nd):
        while nd.left:
            nd=nd.left
        return nd

    def rightmost(self, nd):
        while nd.right:
            nd=nd.right
        return nd

    def find_l(self, v):
        v += 1
        nd = self.root
        prev = 0
        if nd.value < v: prev = nd.value
        while True:
            if v <= nd.value:
                if nd.left:
                    nd = nd.left
                else:
                    return (prev - 1)>>self.id_size
            else:
                prev = nd.value
                if nd.right:
                    nd = nd.right
                else:
                    return (prev - 1)>>self.id_size

    def find_r(self, v):
        v += 1
        nd = self.root
        prev = 0
        if nd.value > v: prev = nd.value
        while True:
            if v < nd.value:
                prev = nd.value
                if nd.left:
                    nd = nd.left
                else:
                    return (prev - 1)>>self.id_size
            else:
                if nd.right:
                    nd = nd.right
                else:
                    return (prev - 1)>>self.id_size

    @property
    def max(self):
        if self.size:
            return self.find_l((1<<(self.N+self.id_size))-1)
        return 0

    @property
    def min(self):
        if self.size:
            return self.find_r(-1)
        return 0

    @property
    def sum(self):
        return self.root.sum-((self.root.value-1)>>self.id_size)

    @property
    def size(self):
        return self.root.cnt-1

    def delete(self,v):
        id=self.val_dic[v]-1
        if id==-1:
            raise IndexError("You are an idiot")
        V=(v<<self.id_size)+id
        self._delete(V)
        self.val_dic[v]-=1

    def _delete(self, v, nd = None, prev = None):
        v += 1
        if not nd: nd = self.root
        if not prev: prev = nd
        while v != nd.value:
            prev = nd
            nd.cnt-=1
            nd.sum-=(v-1)>>self.id_size
            if v <= nd.value:
                if nd.left:
                    nd = nd.left
                else:
                    return
            else:
                if nd.right:
                    nd = nd.right
                else:
                    return
        nd.cnt-=1
        nd.sum-=(v-1)>>self.id_size
        if (not nd.left) and (not nd.right):
            if not prev.left:
                prev.right=None
            elif not prev.right:
                prev.left=None
            else:
                if nd.pivot==prev.left.pivot:
                    prev.left=None
                else:
                    prev.right=None
        elif nd.right:
            nd.value = self.leftmost(nd.right).value
            self._delete(nd.value - 1, nd.right, nd)
        else:
            nd.value = self.rightmost(nd.left).value
            self._delete(nd.value - 1, nd.left, nd)

    def __contains__(self, v: int) -> bool:
        return self.find_r(v - 1) == v

    def __len__(self):
        return self.size

    class node:
        def __init__(self, v, p, m):
            self.value = v
            self.pivot = p
            self.left = None
            self.right = None
            self.cnt=1
            self.sum=(v-1)>>m

N,K = map(int,input().split())
A = list(map(int,input().split()))

if K==1:
    print(A[0])
    exit()

ans = 10**18
BT = MultiBalancingTree(30,id=18)
for i in range(1,N+1):
    if i*K>N:
        break
    tmp = A[i*K-1]
    if i!=1:
        for j in range((i-1)*K,i*K-1):
            BT.append(A[j])
    if i>=2:
        tmp += BT.kth_sum(i-2)[0]
    ans = min(ans,tmp)

print(ans)

yukicoder

結果

テストケース

ソースコード