Problem p02537 (python)

Fixed submission:
import sys
if sys.argv[-1] == 'ONLINE_JUDGE':
    import os
    import re
    with open(__file__) as f:
        source = f.read().split('###''nbacl')
    for s in source[1:]:
        s = re.sub("'''.*", '', s)
        sp = s.split(maxsplit=1)
        if os.path.dirname(sp[0]):
            os.makedirs(os.path.dirname(sp[0]), exist_ok=True)
        with open(sp[0], 'w') as f:
            f.write(sp[1])
import numpy as np
from numpy import int64
from numba import njit
from numba.types import i8
import nbacl.segtree as seg


@njit((i8, i8, i8[:]), cache=True)
def solve(N, K, A):
    dp = np.zeros(300_001, dtype=int64)
    d = seg.segtree(dp, max, 0)
    for i in range(N):
        x = seg.prod(d, max, (A[i] - K) if A[i] > K else 0, (A[i] + K + 1) if A[i] + K < 300_000 else 300_001) + 1
        seg.set(d, max, A[i], x)
    return seg.all_prod(d)


f = open(0)
N, K = [int(x) for x in f.readline().split()]
A = np.fromstring(f.read(), dtype=int64, sep=' ')
ans = solve(N, K, A)
print(ans)

'''
###nbacl nbacl/segtree.py
import numpy as np
from numba import njit
from numba.types import i8


@njit
def ceil_pow2(n):
    x = 0
    while (1 << x) < n:
        x += 1
    return x


@njit
def segtree(a, op, e):
    log = ceil_pow2(a.shape[0])
    size = 1 << log
    d = np.empty(size * 2, dtype=a.dtype)
    d[size:size + len(a)] = a
    d[size + len(a):] = e
    # e
    d[0] = e
    for i in range(size - 1, 0, -1):
        update(d, i, op)
    return d


@njit
def update(d, k, op):
    d[k] = op(d[2 * k], d[2 * k + 1])


@njit
def set(d, op, p, x):
    p += d.shape[0] // 2
    d[p] = x
    while p > 1:
        p >>= 1
        update(d, p, op)


@njit
def get(d, p):
    return d[p + d.shape[0] // 2]


@njit
def prod(d, op, left, right):
    size = d.shape[0] // 2
    sml = d[0]
    smr = d[0]
    left += size
    right += size

    while left < right:
        if left & 1:
            sml = op(sml, d[left])
            left += 1
        if right & 1:
            right -= 1
            smr = op(d[right], smr)
        left >>= 1
        right >>= 1

    return op(sml, smr)


@njit
def all_prod(d):
    return d[1]

'''

Buggy submission (verdict: Runtime Error):
import sys
if sys.argv[-1] == 'ONLINE_JUDGE':
    import os
    import re
    with open(__file__) as f:
        source = f.read().split('###''nbacl')
    for s in source[1:]:
        s = re.sub("'''.*", '', s)
        sp = s.split(maxsplit=1)
        if os.path.dirname(sp[0]):
            os.makedirs(os.path.dirname(sp[0]), exist_ok=True)
        with open(sp[0], 'w') as f:
            f.write(sp[1])
import numpy as np
from numpy import int64
from numba import njit
from numba.types import i8
import nbacl.segtree as seg


@njit((i8, i8, i8[:]), cache=True)
def solve(N, K, A):
    dp = np.zeros(300_001, dtype=int64)
    d = seg.segtree(dp, max)
    for i in range(N):
        x = seg.prod(d, max, A[i] - K, A[i] + K + 1) + 1
        seg.set(d, max, A[i], x)
    return seg.all_prod(d)


f = open(0)
N, K = [int(x) for x in f.readline().split()]
A = np.fromstring(f.read(), dtype=int64, sep=' ')
ans = solve(N, K, A)
print(ans)

'''
###nbacl nbacl/segtree.py
import numpy as np
from numba import njit
from numba.types import i8


@njit
def ceil_pow2(n):
    x = 0
    while (1 << x) < n:
        x += 1
    return x


@njit
def segtree(a, op):
    log = ceil_pow2(a.shape[0])
    size = 1 << log
    d = np.empty(size * 2, dtype=a.dtype)
    d[size:size + len(a)] = a
    # e
    d[0] = 0
    for i in range(size - 1, 0, -1):
        update(d, i, op)
    return d


@njit
def update(d, k, op):
    d[k] = op(d[2 * k], d[2 * k + 1])


@njit
def set(d, op, p, x):
    p += d.shape[0] // 2
    d[p] = x
    while p > 1:
        p >>= 1
        update(d, p, op)


@njit
def get(d, p):
    return d[p + d.shape[0] // 2]


@njit
def prod(d, op, left, right):
    size = d.shape[0] // 2
    sml = d[0]
    smr = d[0]
    left += size
    right += size

    while left < right:
        if left & 1:
            sml = op(sml, d[left])
            left += 1
        if right & 1:
            right -= 1
            smr = op(d[right], smr)
        left >>= 1
        right >>= 1

    return op(sml, smr)


@njit
def all_prod(d):
    return d[1]

'''