Problem p02537 (python)

Fixed submission:
class SegmentTree():
    """
    update, get を提供するSegmentTree

    Attributes
    ----------
    __n : int
        葉の数。2 ^ i - 1
    __dot :
        Segment function
    __e: int
        単位元
    __node: list
        Segment Tree
    """
    def __init__(self, A, dot, e):
        """
        Parameters
        ----------
        A : list
            対象の配列
        dot :
            Segment function
        e : int
            単位元
        """
        n = 2 ** (len(A) - 1).bit_length()
        self.__n = n
        self.__dot = dot
        self.__e = e
        self.__node = [e] * (2 * n)
        for i in range(len(A)):
            self.__node[i + n] = A[i]
        for i in range(n - 1, 0, -1):
            self.__node[i] = self.__dot(self.__node[2 * i], self.__node[2 * i + 1])
    
    def update(self, i, c):
        i += self.__n
        node = self.__node
        node[i] = c
        while i > 1:
            i //= 2
            node[i] = self.__dot(node[2 * i], node[2 * i + 1])

    def get(self, l, r):
        vl, vr = self.__e, self.__e
        l += self.__n
        r += self.__n
        while (l < r):
            if l & 1:
                vl = self.__dot(vl, self.__node[l])
                l += 1
            l //= 2
            if r & 1:
                r -= 1
                vr = self.__dot(vr, self.__node[r])
            r //= 2
        return self.__dot(vl, vr)


N, K = map(int, input().split())
A = [int(input()) for _ in range(N)]
max_A = max(A)

st = SegmentTree([N] * (max_A + 1), min, N)
edge = [set() for _ in range(N)]
for i, a in enumerate(reversed(A)):
    j = N - 1 - i
    p1 = st.get(a, min(max_A + 1, a + K + 1))
    if p1 < N:
        edge[j].add(p1)
    p2 = st.get(max(0, a - K), a)
    if p2 < N:
        edge[j].add(p2)
    p0 = st.get(a, a + 1)
    if p0 < N:
        edge[j].add(p0)
    st.update(a, j)

dp = [1] * N
for i in range(N - 1):
    here = dp[i]
    for ni in edge[i]:
        if ni < N:
            dp[ni] = max(dp[ni], dp[i] + 1)

print(max(dp))

Buggy submission (verdict: Runtime Error):
class SegmentTree():
    """
    update, get を提供するSegmentTree

    Attributes
    ----------
    __n : int
        葉の数。2 ^ i - 1
    __dot :
        Segment function
    __e: int
        単位元
    __node: list
        Segment Tree
    """
    def __init__(self, A, dot, e):
        """
        Parameters
        ----------
        A : list
            対象の配列
        dot :
            Segment function
        e : int
            単位元
        """
        n = 2 ** (len(A) - 1).bit_length()
        self.__n = n
        self.__dot = dot
        self.__e = e
        self.__node = [e] * (2 * n)
        for i in range(len(A)):
            self.__node[i + n] = A[i]
        for i in range(n - 1, 0, -1):
            self.__node[i] = self.__dot(self.__node[2 * i], self.__node[2 * i + 1])
    
    def update(self, i, c):
        i += self.__n
        node = self.__node
        node[i] = c
        while i > 1:
            i //= 2
            node[i] = self.__dot(node[2 * i], node[2 * i + 1])

    def get(self, l, r):
        vl, vr = self.__e, self.__e
        l += self.__n
        r += self.__n
        while (l < r):
            if l & 1:
                vl = self.__dot(vl, self.__node[l])
                l += 1
            l //= 2
            if r & 1:
                r -= 1
                vr = self.__dot(vr, self.__node[r])
            r //= 2
        return self.__dot(vl, vr)


N, K = map(int, input().split())
A = [int(input()) for _ in range(N)]

st = SegmentTree([N] * (max(A) + 1), min, N)
edge = [set() for _ in range(N)]
for i, a in enumerate(reversed(A)):
    j = N - 1 - i
    p1 = st.get(a + 1, a + K + 1)
    if p1 != N:
        edge[j].add(p1)
    p2 = st.get(max(0, a - K), a)
    if p2 != N:
        edge[j].add(p2)
    p0 = st.get(a, a + 1)
    if p0 != j and p0 != N:
        edge[j].add(p0)
    st.update(a, j)

dp = [0] * N
dp[0] = 1
for i in range(N - 1):
    here = dp[i]
    for ni in edge[i]:
        if ni < N:
            dp[ni] = max(dp[ni], dp[i] + 1)

print(max(dp))