AtCoder Beginner Contest 426 E,F,G問題メモ

from math import sqrt


def distance_point_to_line(px, py, sx, sy, tx, ty):
    dx1 = tx - sx
    dy1 = ty - sy
    dx2 = px - sx
    dy2 = py - sy
    t = (dx1 * dx2 + dy1 * dy2) / (dx1 * dx1 + dy1 * dy1)
    if t < 0:
        cx, cy = sx, sy
    elif 1 < t:
        cx, cy = tx, ty
    else:
        cx = sx + t * dx1
        cy = sy + t * dy1
    return sqrt((px - cx) ** 2 + (py - cy) ** 2)


def solve(asx, asy, agx, agy, bsx, bsy, bgx, bgy):
    adx = agx - asx
    ady = agy - asy
    a_dist = sqrt(adx * adx + ady * ady)
    apx = adx / a_dist
    apy = ady / a_dist

    bdx = bgx - bsx
    bdy = bgy - bsy
    b_dist = sqrt(bdx * bdx + bdy * bdy)
    bpx = bdx / b_dist
    bpy = bdy / b_dist

    def distance(t):
        if t > a_dist:
            ax, ay = agx, agy
        else:
            ax = asx + apx * t
            ay = asy + apy * t
        if t > b_dist:
            bx, by = bgx, bgy
        else:
            bx = bsx + bpx * t
            by = bsy + bpy * t
        return sqrt((bx - ax) ** 2 + (by - ay) ** 2)

    l = 0.
    r = min(a_dist, b_dist)
    eps = 1e-7
    while l + eps < r:
        d = (r - l) / 3
        m1 = l + d
        m2 = l + d * 2
        d1 = distance(m1)
        d2 = distance(m2)
        if d1 < d2:
            r = m2
        else:
            l = m1

    ans = distance(l)

    if a_dist < b_dist:
        ans = min(ans, distance_point_to_line(agx, agy, bsx + bpx * a_dist, bsy + bpy * a_dist, bgx, bgy))
    elif a_dist > b_dist:
        ans = min(ans, distance_point_to_line(bgx, bgy, asx + apx * b_dist, asy + apy * b_dist, agx, agy))

    return ans


t = int(input())
for _ in range(t):
    asx, asy, agx, agy = map(int, input().split())
    bsx, bsy, bgx, bgy = map(int, input().split())
    ans = solve(asx, asy, agx, agy, bsx, bsy, bgx, bgy)
    print(ans)

なんか、PyPyで実装するとやたら遅くて、Numbaなら高速に通った。
把握できてない、何らかの遅くなる要素を踏んでいるっぽい。

import os
import sys

import numpy as np


def solve(inp):
    n = inp[0]
    aaa = inp[1:1 + n]
    q = inp[1 + n]
    lll = inp[2 + n::3] - 1
    rrr = inp[3 + n::3]
    kkk = inp[4 + n::3]
    INF = 1 << 60

    # == LazySegmentTree ==
    # Numbaで使うにあたり、いくつかのアドホックな使用方法に注意する。
    # ・data, lazyは2次元配列 (index, 1要素当たりのデータ数) で、Numbaで扱える型。それぞれで型は全て揃える必要がある。
    # ・関数定義(prefixの"lst_"は省略)
    #     operation(i,j,k):  data[k] に data[i]+data[j] を上書き更新
    #     mapping(i,j)    :  data[i] に lazy[j] を上書き反映
    #     composition(i,j):  lazy[i] に lazy[j] を上書き反映
    #     initialize_data(i): data[i] を単位元に戻す
    #     initialize_lazy(i): lazy[i] を単位元に戻す
    # ・初期配列があるときでも、自動的に構築はされないので、自前で前計算しておく
    # ・range_update(l,r) の前には、lazy[0] に更新用の値を埋め込んでおいてから呼ぶ

    lst_n = n  # CUSTOMIZABLE
    lst_depth = 0
    _n = lst_n
    while _n:
        _n >>= 1
        lst_depth += 1
    lst_offset = 1 << lst_depth

    # CUSTOMIZABLE
    lst_data = np.zeros(((lst_offset << 1) + 2, 3), np.int64)
    lst_lazy = np.zeros(lst_offset, np.int64)

    def lst_operation(i, j, k):
        lst_data[k, 0] = lst_data[i, 0] + lst_data[j, 0]
        if lst_data[i, 1] <= lst_data[j, 1]:
            lst_data[k, 1] = lst_data[i, 1]
            lst_data[k, 2] = lst_data[i, 2]
        else:
            lst_data[k, 1] = lst_data[j, 1]
            lst_data[k, 2] = lst_data[j, 2]

    def lst_mapping(i, j):
        lst_data[i, 1] -= lst_lazy[j]

    def lst_composition(i, j):
        lst_lazy[i] += lst_lazy[j]

    def lst_initialize_data(i):
        lst_data[i, 0] = 0
        lst_data[i, 1] = lst_data[i, 2] = INF

    def lst_initialize_lazy(i):
        lst_lazy[i] = 0

    # /CUSTOMIZABLE

    def lst_push(i):
        if i < lst_offset:
            lch = i << 1
            rch = lch + 1
            lst_mapping(lch, i)
            lst_mapping(rch, i)
            if lch < lst_offset:
                lst_composition(lch, i)
                lst_composition(rch, i)
            lst_initialize_lazy(i)

    def lst_update(i):
        lch = i << 1
        rch = lch + 1
        lst_operation(lch, rch, i)

    def lst_all_apply(i):
        lst_mapping(i, 0)
        if i < lst_offset:
            lst_composition(i, 0)

    def lst_propagate(l, r):
        for i in range(lst_depth, 0, -1):
            if ((l >> i) << i) != l:
                lst_push(l >> i)
            if ((r >> i) << i) != r:
                lst_push((r - 1) >> i)

    def lst_range_update(l, r):
        """ lst_lazy[0] に更新したい値を入れてから呼ぶ """
        l += lst_offset
        r += lst_offset
        lst_propagate(l, r)

        l2 = l
        r2 = r
        while l < r:
            if (l & 1) == 1:
                lst_all_apply(l)
                l += 1
            if (r & 1) == 1:
                r -= 1
                lst_all_apply(r)
            l >>= 1
            r >>= 1
        l = l2
        r = r2

        for i in range(1, lst_depth + 1):
            if ((l >> i) << i) != l:
                lst_update(l >> i)
            if ((r >> i) << i) != r:
                lst_update((r - 1) >> i)

    def lst_range_query(l, r):
        l += lst_offset
        r += lst_offset

        lst_propagate(l, r)
        sml = lst_offset << 1
        smr = sml + 1
        lst_initialize_data(sml)
        lst_initialize_data(smr)

        while l < r:
            if (l & 1) == 1:
                lst_operation(sml, l, sml)
                l += 1
            if (r & 1) == 1:
                r -= 1
                lst_operation(r, smr, smr)
            l >>= 1
            r >>= 1
        lst_operation(sml, smr, sml)
        return lst_data[sml]

    def lst_point_set(p, d):
        p += lst_offset
        for i in range(lst_depth, 0, -1):
            lst_push(p >> i)
        lst_data[p] = d
        for i in range(1, lst_depth + 1):
            lst_update(p >> i)

    def lst_point_get(p):
        p += lst_offset
        for i in range(lst_depth, 0, -1):
            lst_push(p >> i)
        return lst_data[p]

    def lst_debug_print():
        i = 1
        while i <= lst_offset:
            print(lst_data[i:i * 2])
            i <<= 1
        i = 1
        while i <= lst_offset:
            print(lst_lazy[i:i * 2])
            i <<= 1

    for i in range(n):
        i2 = i + lst_offset
        lst_data[i2, 0] = 1
        lst_data[i2, 1] = aaa[i]
        lst_data[i2, 2] = i

    for i in range(n, lst_offset):
        i2 = i + lst_offset
        lst_data[i2, 1] = lst_data[i2, 2] = INF

    for i in range(lst_offset - 1, 0, -1):
        lch = i << 1
        rch = lch + 1
        lst_operation(lch, rch, i)

    buf = np.zeros(q, np.int64)
    for qi in range(q):
        l = lll[qi]
        r = rrr[qi]
        k = kkk[qi]
        tmp = 0
        while k > 0:
            cnt, mn, mi = lst_range_query(l, r)
            if cnt == 0:
                break
            if mn > k:
                tmp += cnt * k
                lst_lazy[0] = k
                lst_range_update(l, r)
                break
            tmp += cnt * mn
            k -= mn
            lst_lazy[0] = mn
            lst_range_update(l, r)
            lst_point_set(mi, (0, INF, INF))
            while True:
                _, mn2, mi2 = lst_range_query(l, r)
                if mn2 > 0:
                    break
                lst_point_set(mi2, (0, INF, INF))
        buf[qi] = tmp
    return buf


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('\n'.join(map(str, ans)))

from collections import defaultdict


def solve(n, items, q, queries):
    buf = [0] * q
    classified = defaultdict(list)
    for qi, (l, r, c) in enumerate(queries):
        l -= 1
        r -= 1
        if l == r:
            buf[qi] = items[l][1] if items[l][0] <= c else 0
            continue

        x = l ^ r
        bl = x.bit_length()
        classified[bl, l >> bl].append((qi, l, r, c))

    for (bl, bi), q_list in classified.items():
        w = 1 << bl
        h = w >> 1
        m = bi * w + h
        use_left = set()
        use_right = defaultdict(list)
        max_c = 0
        for qi, l, r, c in q_list:
            use_left.add(l)
            use_right[r].append((qi, l, c))
            max_c = max(max_c, c)

        dp = [0] * (max_c + 1)
        left_cache = {}
        for i in range(m - 1, min(use_left) - 1, -1):
            w, v = items[i]
            for j in range(max_c - w, -1, -1):
                dp[j + w] = max(dp[j + w], dp[j] + v)
            if i in use_left:
                left_cache[i] = dp.copy()

        dp = [0] * (max_c + 1)
        for i in range(m, max(use_right) + 1):
            w, v = items[i]
            for j in range(max_c - w, -1, -1):
                dp[j + w] = max(dp[j + w], dp[j] + v)
            if i in use_right:
                for qi, l, c in use_right[i]:
                    dp_l = left_cache[l]
                    ans = 0
                    for k in range(c + 1):
                        ans = max(ans, dp[k] + dp_l[c - k])
                    buf[qi] = ans

    return buf


n = int(input())
items = [tuple(map(int, input().split())) for _ in range(n)]
q = int(input())
queries = [tuple(map(int, input().split())) for _ in range(q)]
ans = solve(n, items, q, queries)
print('\n'.join(map(str, ans)))