Problem p03153 (python)

Fixed submission:
def main():
    import sys
    from collections import defaultdict
    input = sys.stdin.readline

    class UnionFind():
        def __init__(self, n):
            self.n = n
            self.root = [-1] * (n + 1)
            self.rnk = [0] * (n + 1)

        def find_root(self, x):
            if self.root[x] < 0:
                return x
            else:
                self.root[x] = self.find_root(self.root[x])
                return self.root[x]

        def unite(self, x, y):
            x = self.find_root(x)
            y = self.find_root(y)
            if x == y:
                return
            elif self.rnk[x] > self.rnk[y]:
                self.root[x] += self.root[y]
                self.root[y] = x
            else:
                self.root[y] += self.root[x]
                self.root[x] = y
                if self.rnk[x] == self.rnk[y]:
                    self.rnk[y] += 1

        def isSameGroup(self, x, y):
            return self.find_root(x) == self.find_root(y)

        def size(self, x):
            return -self.root[self.find_root(x)]

    class Bit:
        def __init__(self, n):
            self.size = n
            self.inf = 10 ** 15
            self.tree = [self.inf] * (n + 1)

        def query(self, i):
            s = self.inf
            while i > 0:
                s = min(self.tree[i], s)
                i -= i & -i
            return s

        def update(self, i, x):
            while i <= self.size:
                self.tree[i] = min(self.tree[i], x)
                i += i & -i

    N, D = map(int, input().split())
    A = list(map(int, input().split()))

    tmp = [(a, i+1) for i, a in enumerate(A)]
    tmp.sort(key=lambda x: x[0])
    segtree_plus = Bit(N+5)
    segtree_minus = Bit(N+5)
    edge = []
    val2idx_plus = defaultdict(int)
    val2idx_minus = defaultdict(lambda: N+3)
    edge_append = edge.append
    inf = 10**15
    for a, i in tmp:
        val_plus = a+D*i
        val_minus = a-D*i
        segtree_plus.update(N+1-i, val_plus)
        segtree_minus.update(i, val_minus)
        i_ori = val2idx_plus[val_plus]
        if i > i_ori:
            val2idx_plus[val_plus] = i
        i_ori = val2idx_minus[val_minus]
        if i < i_ori:
            val2idx_minus[val_minus] = i
        plus_min = segtree_plus.query(N-i)
        if plus_min < inf:
            j_plus = val2idx_plus[plus_min]
            edge_append((i, j_plus, A[i-1] + A[j_plus-1] + D*(j_plus - i)))
        minus_min = segtree_minus.query(i-1)
        if minus_min < inf:
            j_minus = val2idx_minus[minus_min]
            edge_append((i, j_minus, A[i-1] + A[j_minus-1] + D * (i - j_minus)))

    edge.sort(key=lambda x: x[2])
    UF = UnionFind(N + 1)
    ans = 0
    for u, v, cost in edge:
        if UF.isSameGroup(u, v):
            continue
        ans += cost
        UF.unite(u, v)
    print(ans)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
def main():
    import sys
    from operator import itemgetter
    input = sys.stdin.buffer.readline

    class UnionFind():
        def __init__(self, n):
            self.n = n
            self.root = [-1] * (n + 1)
            self.rnk = [0] * (n + 1)

        def find_root(self, x):
            if self.root[x] < 0:
                return x
            else:
                self.root[x] = self.find_root(self.root[x])
                return self.root[x]

        def unite(self, x, y):
            x = self.find_root(x)
            y = self.find_root(y)
            if x == y:
                return
            elif self.rnk[x] > self.rnk[y]:
                self.root[x] += self.root[y]
                self.root[y] = x
            else:
                self.root[y] += self.root[x]
                self.root[x] = y
                if self.rnk[x] == self.rnk[y]:
                    self.rnk[y] += 1

        def isSameGroup(self, x, y):
            return self.find_root(x) == self.find_root(y)

        def size(self, x):
            return -self.root[self.find_root(x)]

    class SegTree():
        def __init__(self, N):
            # N:  処理する区間の長さ
            self.N0 = 2 ** (N - 1).bit_length()
            self.INF = 10**15
            self.seg_min = [self.INF] * (2 * self.N0)

        def update(self, index, value):
            index += self.N0 - 1
            self.seg_min[index] = value
            while index:
                index = (index - 1) >> 1
                L = self.seg_min[index * 2 + 1]
                R = self.seg_min[index * 2 + 2]
                if L < R:
                    self.seg_min[index] = L
                else:
                    self.seg_min[index] = R
                #self.seg_min[index] = min(self.seg_min[index * 2 + 1], self.seg_min[index * 2 + 2])

        def query(self, first, last):
            first += self.N0 - 1
            last += self.N0 - 1
            ret = self.INF
            while first <= last:
                if not first & 1:
                    ret_new = self.seg_min[first]
                    if ret_new < ret:
                        ret = ret_new
                    #ret = min(ret, self.seg_min[first])
                if last & 1:
                    ret_new = self.seg_min[last]
                    if ret_new < ret:
                        ret = ret_new
                    #ret = min(ret, self.seg_min[last])
                first = first >> 1
                last = (last >> 1) - 1
            return ret

    N, D = map(int, input().split())
    A = list(map(int, input().split()))

    tmp = [(a, i+1) for i, a in enumerate(A)]
    #tmp.sort(key=lambda x: x[0])
    tmp.sort(key=itemgetter(0))
    segtree_plus = SegTree(N+5)
    segtree_minus = SegTree(N+5)
    edge = []
    val2idx_plus = {}
    val2idx_minus = {}
    edge_append = edge.append
    inf = 10**15
    for a, i in tmp:
        val_plus = a+D*i
        val_minus = a-D*i
        segtree_plus.update(i, val_plus)
        segtree_minus.update(i, val_minus)
        if val_plus not in val2idx_plus:
            val2idx_plus[val_plus] = i
        else:
            i_ori = val2idx_plus[val_plus]
            if i > i_ori:
                val2idx_plus[val_plus] = i
        if val_minus not in val2idx_minus:
            val2idx_minus[val_minus] = i
        else:
            i_ori = val2idx_minus[val_minus]
            if i < i_ori:
                val2idx_minus[val_minus] = i
        plus_min = segtree_plus.query(i+1, N+2)
        if plus_min < inf:
            j_plus = val2idx_plus[plus_min]
            edge_append((i, j_plus, A[i-1] + A[j_plus-1] + D*(j_plus - i)))
        minus_min = segtree_minus.query(1, i-1)
        if minus_min < inf:
            j_minus = val2idx_minus[minus_min]
            edge_append((i, j_minus, A[i-1] + A[j_minus-1] + D * (i - j_minus)))

    #edge.sort(key=lambda x: x[2])
    edge.sort(key=itemgetter(2))
    UF = UnionFind(N + 1)
    ans = 0
    for u, v, cost in edge:
        if UF.isSameGroup(u, v):
            continue
        ans += cost
        UF.unite(u, v)
    print ans


if __name__ == '__main__':
    main()