AtCoder Beginner Contest 252 G,Ex問題メモ

import os
import sys

import numpy as np


def solve(inp):
    n = inp[0]
    ppp = inp[1:]
    MOD = 998244353

    dp = np.zeros((n + 1, n + 1), dtype=np.int64)  # [l,r) から構築できる木
    for i in range(n):
        dp[i, i] = dp[i, i + 1] = 1

    for w in range(2, n + 1):
        for l in range(n - w + 1):
            r = l + w
            tmp = dp[l + 1, r]
            for k in range(l + 1, r):
                if ppp[l] < ppp[k]:
                    tmp += dp[l + 1, k] * dp[k, r]
                    tmp %= MOD
            dp[l, r] = tmp

    return dp[1, n]


SIGNATURE = '(i8[:],)'
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve', SIGNATURE)(solve)
    cc.compile()
    exit()

if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit(SIGNATURE, cache=True)(solve)
    print('compiled', file=sys.stderr)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)

import os
import sys

import numpy as np


def solve(inp):
    n, c, k = inp[:3]
    ddd = inp[3::2] - 1
    vvv = inp[4::2]

    if c == 1:
        vvv = np.sort(vvv)[::-1]
        return vvv[k - 1]

    int_list = [n]
    int_list.clear()
    gems = [int_list.copy() for _ in range(c)]

    for i in range(n):
        d = ddd[i]
        v = vvv[i]
        gems[d].append(v)

    gems_counts = np.array([len(g) for g in gems], np.int64)

    all_prod = 1
    can_prods = [{np.int64(1)}]
    for gc in gems_counts:
        all_prod *= gc
        can_prod = can_prods[-1].copy()
        for p in can_prods[-1]:
            can_prod.add(p * gc)
        can_prods.append(can_prod)

    can_prod_s = sorted(can_prods[-1])
    tgt = can_prod_s[len(can_prod_s) // 2]

    group0 = []
    group1 = []
    for i in range(c - 1, -1, -1):
        g = gems[i]
        gc = gems_counts[i]
        if tgt % gc == 0 and tgt // gc in can_prods[i]:
            group0.append(g)
            tgt //= gc
        else:
            group1.append(g)

    def all_pattern(group):
        res = [0]
        for g in group:
            new_res = []
            for v in g:
                for u in res:
                    new_res.append(v ^ u)
            res = new_res
        return res

    pat0 = all_pattern(group0)
    pat1 = all_pattern(group1)

    lll0 = [-1]
    rrr0 = [-1]
    prp0 = [False]
    num0 = [pat0]
    cnt0 = [len(pat0)]

    lll1 = [-1]
    rrr1 = [-1]
    prp1 = [False]
    num1 = [pat1]
    cnt1 = [len(pat1)]

    def propagate(lll, rrr, prp, num, cnt, i, d):
        l_num = []
        r_num = []
        for v in num[i]:
            if v & d:
                r_num.append(v)
            else:
                l_num.append(v)
        if len(l_num) > 0:
            lll[i] = len(lll)
            lll.append(-1)
            rrr.append(-1)
            prp.append(False)
            num.append(l_num)
            cnt.append(len(l_num))
        if len(r_num) > 0:
            rrr[i] = len(lll)
            lll.append(-1)
            rrr.append(-1)
            prp.append(False)
            num.append(r_num)
            cnt.append(len(r_num))
        prp[i] = True

    def check(m):
        q = [(0, 0, 1 << 59)]
        res = 0
        while q:
            i0, i1, d = q.pop()
            if d == 0:
                res += cnt0[i0] * cnt1[i1]
                continue

            if cnt0[i0] * cnt1[i1] < 1000:
                for v0 in num0[i0]:
                    for v1 in num1[i1]:
                        if v0 ^ v1 >= m:
                            res += 1
                continue

            if not prp0[i0]:
                propagate(lll0, rrr0, prp0, num0, cnt0, i0, d)

            if not prp1[i1]:
                propagate(lll1, rrr1, prp1, num1, cnt1, i1, d)

            if d & m:
                if lll0[i0] != -1 and rrr1[i1] != -1:
                    q.append((lll0[i0], rrr1[i1], d >> 1))
                if rrr0[i0] != -1 and lll1[i1] != -1:
                    q.append((rrr0[i0], lll1[i1], d >> 1))
            else:
                if lll0[i0] != -1 and rrr1[i1] != -1:
                    res += cnt0[lll0[i0]] * cnt1[rrr1[i1]]
                if rrr0[i0] != -1 and lll1[i1] != -1:
                    res += cnt0[rrr0[i0]] * cnt1[lll1[i1]]
                if lll0[i0] != -1 and lll1[i1] != -1:
                    q.append((lll0[i0], lll1[i1], d >> 1))
                if rrr0[i0] != -1 and rrr1[i1] != -1:
                    q.append((rrr0[i0], rrr1[i1], d >> 1))

            if res >= k:
                return k

            d >>= 1

        return res

    l = 0
    r = 1 << 60
    while l + 1 < r:
        m = (l + r) // 2
        if check(m) < k:
            r = m
        else:
            l = m
    return l


SIGNATURE = '(i8[:],)'
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve', SIGNATURE)(solve)
    cc.compile()
    exit()

if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit(SIGNATURE, cache=True)(solve)
    print('compiled', file=sys.stderr)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)