#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
コンパイル時間	818 ms
コンパイル使用メモリ	86,876 KB
実行使用メモリ	136,960 KB
最終ジャッジ日時	2023-09-24 12:25:59
合計ジャッジ時間	12,688 ms
ジャッジサーバーID （参考情報）	judge15 / judge12

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

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	75 ms 71,128 KB
testcase_01	AC	75 ms 71,320 KB
testcase_02	AC	75 ms 71,304 KB
testcase_03	AC	75 ms 71,056 KB
testcase_04	AC	74 ms 71,128 KB
testcase_05	WA	-
testcase_06	WA	-
testcase_07	AC	74 ms 71,272 KB
testcase_08	WA	-
testcase_09	WA	-
testcase_10	WA	-
testcase_11	WA	-
testcase_12	WA	-
testcase_13	AC	132 ms 101,592 KB
testcase_14	WA	-
testcase_15	AC	124 ms 101,720 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	73 ms 71,164 KB
testcase_24	WA	-
testcase_25	AC	319 ms 112,104 KB
testcase_26	WA	-
testcase_27	AC	128 ms 101,536 KB
testcase_28	AC	128 ms 101,696 KB
testcase_29	AC	126 ms 102,112 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

結果

テストケース

ソースコード