Submission #3441934


Source Code Expand

def rotate(nd, d):
    c = nd[d]
    if d:
        e = c[1]
        nd[1] = c[0]
        c[0] = nd
    else:
        e = c[0]
        nd[0] = c[1]
        c[1] = nd

    r = c[3] = nd[3]
    nd[3] = r - (e[3] if e else 0) - 1
    return c

def insert(t, val):
    x = t
    st = []; dr = []
    y = None
    while x:
        y = x
        st.append(x)
        if x[2] == val:
            return t
        d = (x[2] < val)
        dr.append(d)
        x = x[d]

    # [<left>, <right>, <key>, <count>, <color: [BLACK,RED]>]
    nd = [None, None, val, 1, 1]

    if not y:
        nd[4] = 0
        return nd

    y[d] = nd
    for x in st:
        x[3] += 1

    while 1:
        if not st:
            # if nd is the root node
            nd[4] = 0
            return nd

        # nd's parent
        y = st.pop(); dy = dr.pop()
        if not y[4]:
            # if nd's parent is black
            return t

        # nd's grandparent
        w = st.pop(); dw = dr.pop()
        # nd's uncle
        z = w[dw ^ 1]
        if z and z[4]:
            # the parent and the uncle is red
            y[4] = z[4] = 0
            w[4] = 1
            nd = w
            continue

        if dw != dy:
            w[dw] = rotate(y, dy)
            nd, y = y, nd

        rotate(w, dw)
        y[4] = 0
        w[4] = 1

        if st:
            st[-1][dr[-1]] = y
            return t
        return y

def __delete(st, dr, nd):
    x = nd
    if x[0]:
        st.append(x); dr.append(0)
        x = x[0]
        while x:
            st.append(x); dr.append(1)
            x = x[1]

        y = st.pop(); dr.pop()
        # copy a value only
        nd[2] = y[2]
    elif x[1]:
        st.append(x); dr.append(1)
        x = x[1]
        while x:
            st.append(x); dr.append(0)
            x = x[0]

        y = st.pop(); dr.pop()
        # copy a value only
        nd[2] = y[2]
    else:
        y = nd

    #assert len(st) == len(dr)

    # delete the node (y) with at most one non-leaf child
    for x in st:
        x[3] -= 1

    c = y[0] if y[0] else y[1]
    if st:
        st[-1][dr[-1]] = c

    if y[4]:
        # if nd is red
        return st[0] if st else c

    if c and c[4]:
        # if nd is black and child is red
        c[4] = 0
        return st[0] if st else c

    # nd and child is black: delete in each cases

    n = c
    while st:
        p = st.pop(); dp = dr.pop()
        # n's sibling
        s = p[dp ^ 1]
        #assert s
        sc = s[dp]; sd = s[dp ^ 1]

        if s[4]:
            #assert sc and sd
            p[4] = 1; s[4] = 0
            if st:
                st[-1][dr[-1]] = rotate(p, dp ^ 1)
            else:
                rotate(p, dp ^ 1)
            st.append(s); dr.append(dp)
            #assert p[dp ^ 1] is sc
            s = sc
            sc = s[dp]; sd = s[dp ^ 1]
            break

        if p[4] or (sc and sc[4]) or (sd and sd[4]):
            break

        s[4] = 1
        n = p
    else:
        return n

    if p[4] and not s[4] and not (sc and sc[4]) and not (sd and sd[4]):
        p[4] = 0; s[4] = 1
        return st[0] if st else p

    sc = s[dp]; sd = s[dp ^ 1]
    if not s[4] and (sc and sc[4]) and not (sd and sd[4]):
        sc[4] = 0; s[4] = 1
        sd = s
        s = p[dp ^ 1] = rotate(s, dp)

    #assert not s[4] and (sd and sd[4])
    s[4] = p[4]
    p[4] = 0
    sd[4] = 0

    if st:
        st[-1][dr[-1]] = rotate(p, dp ^ 1)
    else:
        return rotate(p, dp ^ 1)

    return st[0] if st else n

def delete(t, val):
    x = t

    st = []; dr = []
    while x and val != x[2]:
        d = (x[2] < val)
        st.append(x); dr.append(d)
        x = x[d]
    if not x:
        return t
    # x is the node with key = val
    return __delete(st, dr, x)

root = None
def query(k):
    global root
    x = root
    st = []; dr = []
    while x:
        l = x[0]
        c = (l[3] if l else 0) + 1

        if c == k:
            break

        st.append(x)
        if c < k:
            k -= c
            x = x[1]
            dr.append(1)
        else:
            x = x[0]
            dr.append(0)
    v = x[2]
    root = __delete(st, dr, x)
    return v


import sys
readline = sys.stdin.readline
writelines = sys.stdout.writelines

Q = int(readline())
ans = []
for i in range(Q):
    t, x = map(int, readline().split())
    if t == 1:
        root = insert(root, x)
    else:
        ans.append("%d\n" % query(x))
writelines(ans)

Submission Info

Submission Time
Task C - データ構造
User yaketake08
Language PyPy3 (2.4.0)
Score 100
Code Size 4658 Byte
Status AC
Exec Time 1406 ms
Memory 114780 KB

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 100 / 100
Status
AC × 2
AC × 18
Set Name Test Cases
Sample sample_01.txt, sample_02.txt
All sample_01.txt, sample_02.txt, subtask1_01.txt, subtask1_02.txt, subtask1_03.txt, subtask1_04.txt, subtask1_05.txt, subtask1_06.txt, subtask1_07.txt, subtask1_08.txt, subtask1_09.txt, subtask1_10.txt, subtask1_11.txt, subtask1_12.txt, subtask1_13.txt, subtask1_14.txt, subtask1_15.txt, subtask1_16.txt
Case Name Status Exec Time Memory
sample_01.txt AC 163 ms 38256 KB
sample_02.txt AC 163 ms 38256 KB
subtask1_01.txt AC 165 ms 38256 KB
subtask1_02.txt AC 164 ms 38256 KB
subtask1_03.txt AC 164 ms 38256 KB
subtask1_04.txt AC 392 ms 56156 KB
subtask1_05.txt AC 509 ms 63324 KB
subtask1_06.txt AC 171 ms 38768 KB
subtask1_07.txt AC 875 ms 83420 KB
subtask1_08.txt AC 1061 ms 89180 KB
subtask1_09.txt AC 1063 ms 90076 KB
subtask1_10.txt AC 1266 ms 104284 KB
subtask1_11.txt AC 1244 ms 101596 KB
subtask1_12.txt AC 997 ms 91996 KB
subtask1_13.txt AC 1391 ms 113884 KB
subtask1_14.txt AC 1406 ms 114780 KB
subtask1_15.txt AC 828 ms 79964 KB
subtask1_16.txt AC 1002 ms 90076 KB