Problem:


Score : 600 points

Problem Statement

There are N cities in Republic of AtCoder. The size of the i-th city is A_{i}.
Takahashi would like to build N-1 bidirectional roads connecting two cities so that any city can be reached from any other city by using these roads.

Assume that the cost of building a road connecting the i-th city and the j-th city is |i-j| \times D + A_{i} + A_{j}.
For Takahashi, find the minimum possible total cost to achieve the objective.

Constraints

1 \leq N \leq 2 \times 10^5

1 \leq D \leq 10^9

1 \leq A_{i} \leq 10^9

A_{i} and D are integers.

Input

Input is given from Standard Input in the following format:

N D
A_1 A_2 ... A_N

Output

Print the minimum possible total cost.

Sample Input 1

3 1
1 100 1

Sample Output 1

106

This cost can be achieved by, for example, building roads connecting City 1, 2 and City 1, 3.

Sample Input 2

3 1000
1 100 1

Sample Output 2

2202

Sample Input 3

6 14
25 171 7 1 17 162

Sample Output 3

497

Sample Input 4

12 5
43 94 27 3 69 99 56 25 8 15 46 8

Sample Output 4

658



Sample input:
3 1
1 100 1


Sample output:
106

Submission status: Time Limit Exceeded

Solution code:
###template###
import sys
def input(): return sys.stdin.readline().rstrip()
from collections import defaultdict, Counter
from itertools import product, groupby, count, permutations, combinations
from math import pi, sqrt, ceil, floor
from collections import deque
from bisect import bisect, bisect_left, bisect_right
from string import ascii_lowercase
import heapq
INF = float("inf")
sys.setrecursionlimit(10**7)
# 4近傍（右, 下, 左, 上）
dy = [0, -1, 0, 1]
dx = [1, 0, -1, 0]
def inside(y: int, x: int, H: int, W: int) -> bool: return 0 <= y < H and 0 <= x < W
def mi(): return map(int, input().split())
def ii(): return int(input())
###template###

N, D = mi()

#As[i] = [各都市のコスト,各都市の番号]
#0-indexed
As = [[a, i] for i, a in enumerate(list(mi()))]

#Union-Find ****0スタートなので注意！****
class UnionFind():
  # 作りたい要素数（ノードの数）nで初期化
  # 使用するインスタンス変数の初期化
  def __init__(self, n):
    self.n = n
    # root[x]<0ならそのノードが根かつその値が木の要素数
    # rootノードでその木の要素数を記録する
    self.root = [-1]*(n) #元々はn+1になっていたが、nにした（n+1なら1スタートで使えるかもしれないが、Count_Trees()がおかしくなるのでやめた）
    # 木をくっつける時にアンバランスにならないように調整する
    self.rnk = [0]*(n) #元々はn+1になっていたが、nにした

  # ノードxのrootノードを見つける
  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):
    # 入力ノードのrootノードを見つける
    x = self.Find_Root(x)
    y = self.Find_Root(y)
    # すでに同じ木に属していた場合、何もしないで返す
    if(x == y):
      return
    # 違う木に属していた場合rnkを見てくっつける方を決める
    elif(self.rnk[x] > self.rnk[y]):
      self.root[x] += self.root[y] #要素数を足す(根のrootは"-要素数")
      self.root[y] = x #ノードyの親をx（新しい集合の根）とする
      return 1
    else:
      self.root[y] += self.root[x] #要素数を足す
      self.root[x] = y
      # rnkが同じ（深さに差がない場合）は1増やす（根(rnk1)+繋げる集合のrnkということで、1増えることになる）
      if(self.rnk[x] == self.rnk[y]):
          self.rnk[y] += 1
      return 1

  # xとyが同じグループに属するか判断
  def isSameGroup(self, x, y):
    return self.Find_Root(x) == self.Find_Root(y)

  # ノードxが属する木のサイズを返す（根ノードの.root[idx]を-符号を外して返すだけ
  def Count(self, x):
    return -self.root[self.Find_Root(x)] #-符号を外す

#As_規模降順 = [各都市の規模, 各都市の番号（降順ソート前の番号）]
from operator import itemgetter
As_costdescending = sorted(As, key=itemgetter(0), reverse=True)

#やっぱSegTree使う
seg_length = N
bitlen = (seg_length-1).bit_length()
N0 = 2**bitlen
#INF = 2**31-1 #自分のtemplateで宣言してるのでいらない
data = [[INF,0] for _ in range(2*N0)]
#各ノードにおける最小値のidxの初期値を格納
for i in range(1, bitlen+1): #(0は除き)1,2,3段目に行う iは段数を表す
  for j in range(0,2**i):
    #一番左は0のまま、残りはidxを変える
    #i段目の左端のidxは2**i-1
    data[2**i-1 + j][1] = (2**(bitlen-i))*j
# a_k の値を x に更新
def update(k, x):
  firstk = k
  k += N0-1 #この時のk : セグ木上の最下段のidx
  data[k][0] = x
  while k >= 0:
#    print('最初はkは',firstk, 'でしたが今は', k, x, k+N0-1)
    k = (k - 1) // 2
    left = data[2*k+1][0]
    right = data[2*k+2][0]
    if left <= right: #左側が小さければ親ノードに左に入っているidxを送る
      data[k][1] = data[2*k+1][1]
    else: data[k][1] = data[2*k+2][1] #右側が小さかったら
    data[k][0] = min(left,right)
    #右隣のノードと比べて、上のノードに最小値を返す？
# 区間[l, r)の最小値
def query(l, r):
  L = l + N0; R = r + N0
  s = INF
  minidx = 0
  while L < R:
    if R & 1:
      R -= 1
      if data[R-1][0]<s:
        minidx = data[R-1][1]
        s = data[R-1][0]

    if L & 1:
      if data[L-1][0]<s:
        minidx = data[L-1][1]
        s = data[L-1][0]
      L += 1
    L >>= 1; R >>= 1
  return s, minidx
#SegTreeおわり

#もう一個セグ木をつくる（力技）
data2 = [[INF,0] for _ in range(2*N0)]
#各ノードにおける最小値のidxの初期値を格納
for i in range(1, bitlen+1): #(0は除き)1,2,3段目に行う iは段数を表す
  for j in range(0,2**i):
    #一番左は0のまま、残りはidxを変える
    #i段目の左端のidxは2**i-1
    data2[2**i-1 + j][1] = (2**(bitlen-i))*j
# a_k の値を x に更新
def update_(k, x):
  firstk = k
  k += N0-1 #この時のk : セグ木上の最下段のidx
  data2[k][0] = x
  while k >= 0:
#    print('最初はkは',firstk, 'でしたが今は', k, x, k+N0-1)
    k = (k - 1) // 2
    left = data2[2*k+1][0]
    right = data2[2*k+2][0]
    if left <= right: #左側が小さければ親ノードに左に入っているidxを送る
      data2[k][1] = data2[2*k+1][1]
    else: data2[k][1] = data2[2*k+2][1] #右側が小さかったら
    data2[k][0] = min(left,right)
    #右隣のノードと比べて、上のノードに最小値を返す？
# 区間[l, r)の最小値
def query_(l, r):
  L = l + N0; R = r + N0
  s = INF
  minidx = 0
  while L < R:
    if R & 1:
      R -= 1
      if data2[R-1][0]<s:
        minidx = data2[R-1][1]
        s = data2[R-1][0]

    if L & 1:
      if data2[L-1][0]<s:
        minidx = data2[L-1][1]
        s = data2[L-1][0]
      L += 1
    L >>= 1; R >>= 1
  return s, minidx
#SegTree_おわり

for a, i in As:
  update(i, (N-1-i)*D+a) #sgt_left
  update_(i, i*D+a) #sgt_right


#print(cummin_left)
#print(cummin_right)
#[(cost, idx, adjidx), ...]
edge_candidates = []

for a, o_idx in As_costdescending: #o_idx:降順にする前のidx
  if o_idx == 0: #右だけを見る
    rightmin, rightmin_oidx = query_(1,N) #[L,r)
    edge_candidates.append((abs(o_idx-rightmin_oidx)*D+a+rightmin-rightmin_oidx*D, o_idx, rightmin_oidx))
  elif o_idx == N-1: #左だけを見る
    leftmin, leftmin_oidx = query(0, N-1) #[L,r)
    edge_candidates.append((abs(o_idx-leftmin_oidx)*D+a+leftmin-(N-1-leftmin_oidx)*D, o_idx, leftmin_oidx))
  else:
    rightmin, rightmin_oidx = query_(o_idx+1,N) #[L,r)
    edge_candidates.append((abs(o_idx-rightmin_oidx)*D+a+rightmin-rightmin_oidx*D, o_idx, rightmin_oidx))
    leftmin, leftmin_oidx = query(0, o_idx) #[L,r)
    edge_candidates.append((abs(o_idx-leftmin_oidx)*D+a+leftmin-(N-1-leftmin_oidx)*D, o_idx, leftmin_oidx))
  update(o_idx,INF)
  update_(o_idx,INF)


heapq.heapify(edge_candidates)

#上でできた配列に対してクラスカル法を適用して最小全域木（のコスト総和）を求める
#Union-Findを使ったクラスカル法
uf = UnionFind(N)
ans = 0
while edge_candidates:
  c, idx, adjidx = heapq.heappop(edge_candidates)
  if uf.Unite(idx, adjidx): ans += c

print(ans)

Submission status: Time Limit Exceeded

Solution code:
class SegmentTree:
    def __init__(self, n):
        k = 1
        while k < n:
            k *= 2
        self.N = k
        self.node = [(float("inf"), 0)] * (int(2 * self.N - 1))

    def update(self, i, x):
        i += self.N - 1
        self.node[i] = x
        while i > 0:
            i = (i - 1) // 2
            self.node[i] = min(self.node[2 * i + 1], self.node[2 * i + 2])

    def _find(self, a, b, i, l, r):
        if r <= a or b <= l:
            return float("inf"), 0
        elif a <= l and r <= b:
            return self.node[i]

        else:
            mid = (l + r) // 2
            return min(self._find(a, b, 2 * i + 1, l, mid), self._find(a, b, 2 * i + 2, mid, r))

    def find(self, a, b):
        return self._find(a, b + 1, 0, 0, self.N)


N, D = map(int, input().split())
A = list(map(int, input().split()))
parent = [i for i in range(N)]
rank = [0] * N


def find(i):
    if parent[i] == i:
        return i
    else:
        parent[i] = find(parent[i])
        return parent[i]


def same(x, y):
    return find(x) == find(y)


def unite(x, y):
    x = find(x)
    y = find(y)
    if x == y:
        return

    if rank[x] > rank[y]:
        parent[y] = x
    else:
        parent[x] = y
        if rank[x] == rank[y]:
            rank[y] += 1


INF = float('inf')
R = [i for i in range(N)]
R = sorted(R, key=lambda x: A[x])
B = [0] * N
C = [0] * N
for i in range(N):
    B[i] = A[i] + D * i

for i in range(N):
    C[i] = A[i] + D * (N - i)

stL = SegmentTree(N)
stR = SegmentTree(N)
E = []
for i in R:
    p = stL.find(0, i)
    if p[0] != INF:
        cost = p[0] + A[i] - (N - i) * D
        E.append((cost, i, p[1]))
    p = stR.find(i, N - 1)
    if p[0] != INF:
        cost = p[0] + A[i] - i * D
        E.append((cost, i, p[1]))

    stL.update(i, (A[i] + (N - i) * D, i))
    stR.update(i, (A[i] + i * D, i))

E = sorted(E)
ans = 0
for e in E:
    if not same(e[1], e[2]):
        unite(e[1], e[2])
        ans += e[0]

print(ans)

Submission status: Time Limit Exceeded

Solution code:
# from typing import Callable


def connecting_cities(N: int, D: int, A: list) -> int:
    edges = select_edges(N, D, A)
    uft = UnionFindTree(N)

    count = 0
    for u, v, d in sorted(edges, key=lambda x: x[2]):
        if not uft.same(u, v):
            uft.summarize(u, v)
            count += d

    return count


def select_edges(N: int, D: int, A: list) -> list:
    INF = 1 << 32
    st_l = SegmentTree(N, lambda x, y: x if x[0] < y[0] else y, (INF, -1))
    st_r = SegmentTree(N, lambda x, y: x if x[0] < y[0] else y, (INF, -1))
    indexes_A = sorted([i for i in range(N)], key=lambda x: A[x])

    edges = []
    for i in indexes_A:
        la, li = st_l.range(0, i)
        if la < INF and li != -1:
            edges.append((li, i, abs(i - li) * D + A[i] + A[li]))

        ra, ri = st_r.range(i, N)
        if ra < INF and ri != -1:
            edges.append((ri, i, abs(i - ri) * D + A[i] + A[ri]))

        st_l.update(i, (A[i] + (N - i) * D, i))
        st_r.update(i, (A[i] + i * D, i))

    return edges


class SegmentTree:

    # def __init__(self, size: int, op: Callable[[any, any], any], init: any = 0):
    def __init__(self, size, op, init):
        """initialize SegmentTree
        :param size: Size of tree. This must be natural number.
        :param op: Operator which compares two numbers. This function return the
                representative value of two arguments.
        :param init: The initial value of element of tree.
        """
        if size < 1:
            raise Exception('size must be greater than 0 (actual = %d)' % size)

        self.__treesize = 1
        self.__init = init
        # tree size is the minimum number which is greater than or equal to
        # designed `size` and is power of 2.
        while self.__treesize < size:
            self.__treesize *= 2

        self.__size = self.__treesize
        self.__treesize = self.__treesize * 2 - 1

        self.__comp = op

        # initialize tree
        self.__tree = [init for _ in range(self.__treesize)]

    def range(self, l: int, r: int) -> int:
        """Returns the representative value in range [l, r)
        :param l: left value(include)
        :param r: right value(not include)
        :return: the representative value in range[l, r)
        """
        return self.__range(l, r, 0, 0, self.__size)

    def __range(self, l: int, r: int, k: int, kl: int, kr: int) -> int:
        """
        :param l: left value (include)
        :param r: right value (not include)
        :param k: node index
        :param kl: left value of node
        :param kr: right value of node
        """
        if kr <= l or r <= kl:
            # no crossing
            return self.__init
        if l <= kl and kr <= r:
            # including whole k's range
            return self.__tree[k]

        vl = self.__range(l, r, 2 * k + 1, kl, (kl + kr) // 2)
        vr = self.__range(l, r, 2 * k + 2, (kl + kr) // 2, kr)
        if vl is None:
            return vr
        if vr is None:
            return vl

        return self.__comp(vl, vr)

    def update(self, index: int, val: int):
        """update value at self.__tree[index]'s value
        """
        index += self.__size - 1
        self.__tree[index] = val

        while index > 0:
            index = (index - 1) // 2
            self.__tree[index] = self.__comp(
                self.__tree[2*index + 1], self.__tree[2*index + 2])

    def print(self):
        """for debug
        """
        print(self.__tree)


class UnionFindTree:

    def __init__(self, size: int):
        """UnionFindTreeを初期化します
        :param size: サイズ
        :return: UnionFindTree
        """
        if size < 0:
            raise Exception('size must be greater than 0.')

        self.__parent = [-1] * size

    def summarize(self, a: int, b: int):
        """aを含む木とbを含む木をまとめます
        :param a: 木a
        :param b: 木b
        """
        a = self.__root(a)
        b = self.__root(b)

        if a != b:
            self.__parent[b] = a

    def __root(self, a: int) -> int:
        """aの根を求めます
        :param a: 木の要素
        :return: 根
        """
        if self.__parent[a] == -1:
            return a
        else:
            self.__parent[a] = self.__root(self.__parent[a])
            return self.__parent[a]

    def same(self, a: int, b: int) -> bool:
        """a と b が同じ木に属するかを判断します
        :param a: 木
        :param b: 木
        :return: a と b が同じ木なら True。そうでないなら False。
        """
        return self.__root(a) == self.__root(b)


if __name__ == "__main__":
    N, D = [int(s) for s in input().split()]
    A = [int(s) for s in input().split()]
    ans = connecting_cities(N, D, A)
    print(ans)

Submission status: Time Limit Exceeded

Solution code:
# from typing import Callable


def connecting_cities(N: int, D: int, A: list) -> int:
    edges = select_edges(N, D, A)
    uft = UnionFindTree(N)

    count = 0
    for u, v, d in sorted(edges, key=lambda x: x[2]):
        if not uft.same(u, v):
            uft.summarize(u, v)
            count += d

    return count


def select_edges(N: int, D: int, A: list) -> list:
    INF = 1 << 32
    st_l = SegmentTree(N, lambda x, y: x if x[0] < y[0] else y, (INF, -1))
    st_r = SegmentTree(N, lambda x, y: x if x[0] < y[0] else y, (INF, -1))
    indexes_A = sorted([i for i in range(N)], key=lambda x: A[x])

    edges = []
    for i in indexes_A:
        la, li = st_l.range(0, i)
        if la < INF and li != -1:
            edges.append((li, i, abs(i - li) * D + A[i] + A[li]))

        ra, ri = st_r.range(i, N)
        if ra < INF and ri != -1:
            edges.append((ri, i, abs(i - ri) * D + A[i] + A[ri]))

        st_l.update(i, (A[i] + (N - i) * D, i))
        st_r.update(i, (A[i] + i * D, i))

    return edges


class SegmentTree:

    # def __init__(self, size: int, op: Callable[[any, any], any], init: any = 0):
    def __init__(self, size, op, init):
        """initialize SegmentTree
        :param size: Size of tree. This must be natural number.
        :param op: Operator which compares two numbers. This function return the
                representative value of two arguments.
        :param init: The initial value of element of tree.
        """
        if size < 1:
            raise Exception('size must be greater than 0 (actual = %d)' % size)

        self.__treesize = 1
        self.__init = init
        # tree size is the minimum number which is greater than or equal to
        # designed `size` and is power of 2.
        while self.__treesize < size:
            self.__treesize *= 2

        self.__size = self.__treesize
        self.__treesize = self.__treesize * 2 - 1

        self.__comp = op

        # initialize tree
        self.__tree = [init for _ in range(self.__treesize)]

    def range(self, l: int, r: int) -> int:
        """Returns the representative value in range [l, r)
        :param l: left value(include)
        :param r: right value(not include)
        :return: the representative value in range[l, r)
        """
        return self.__range(l, r, 0, 0, self.__size)

    def __range(self, l: int, r: int, k: int, kl: int, kr: int) -> int:
        """
        :param l: left value (include)
        :param r: right value (not include)
        :param k: node index
        :param kl: left value of node
        :param kr: right value of node
        """
        if kr <= l or r <= kl:
            # no crossing
            return self.__init
        if l <= kl and kr <= r:
            # including whole k's range
            return self.__tree[k]

        vl = self.__range(l, r, 2 * k + 1, kl, (kl + kr) // 2)
        vr = self.__range(l, r, 2 * k + 2, (kl + kr) // 2, kr)
        if vl is None:
            return vr
        if vr is None:
            return vl

        return self.__comp(vl, vr)

    def update(self, index: int, val: int):
        """update value at self.__tree[index]'s value
        """
        index += self.__size - 1
        self.__tree[index] = val

        while index > 0:
            index = (index - 1) // 2
            self.__tree[index] = self.__comp(
                self.__tree[2*index + 1], self.__tree[2*index + 2])

    def print(self):
        """for debug
        """
        print(self.__tree)


class UnionFindTree:

    def __init__(self, size: int):
        """UnionFindTreeを初期化します
        :param size: サイズ
        :return: UnionFindTree
        """
        if size < 0:
            raise Exception('size must be greater than 0.')

        self.__parent = [-1] * size

    def summarize(self, a: int, b: int):
        """aを含む木とbを含む木をまとめます
        :param a: 木a
        :param b: 木b
        """
        a = self.__root(a)
        b = self.__root(b)

        if a != b:
            self.__parent[b] = a

    def __root(self, a: int) -> int:
        """aの根を求めます
        :param a: 木の要素
        :return: 根
        """
        if self.__parent[a] == -1:
            return a
        else:
            return self.__root(self.__parent[a])

    def same(self, a: int, b: int) -> bool:
        """a と b が同じ木に属するかを判断します
        :param a: 木
        :param b: 木
        :return: a と b が同じ木なら True。そうでないなら False。
        """
        return self.__root(a) == self.__root(b)


if __name__ == "__main__":
    N, D = [int(s) for s in input().split()]
    A = [int(s) for s in input().split()]
    ans = connecting_cities(N, D, A)
    print(ans)

Submission status: Time Limit Exceeded

Solution code:
# from typing import Callable


def connecting_cities(N: int, D: int, A: list) -> int:
    edges = select_edges(N, D, A)
    edges = sorted(edges, key=lambda x: x[2], reverse=True)

    uft = UnionFindTree(N)

    count = 0
    while not len(edges) == 0:
        u, v, d = edges.pop()

        if not uft.same(u, v):
            uft.summarize(u, v)
            count += d

    return count


def select_edges(N: int, D: int, A: list) -> list:
    INF = 1 << 32
    st_l = SegmentTree(N, lambda x, y: x if x[0] < y[0] else y, (INF, -1))
    st_r = SegmentTree(N, lambda x, y: x if x[0] < y[0] else y, (INF, -1))
    indexes_A = sorted([i for i in range(N)], key=lambda x: A[x])

    edges = []
    for i in indexes_A:
        la, li = st_l.range(0, i)
        if la < INF and li != -1:
            edges.append((li, i, abs(i - li) * D + A[i] + A[li]))

        ra, ri = st_r.range(i, N)
        if ra < INF and ri != -1:
            edges.append((ri, i, abs(i - ri) * D + A[i] + A[ri]))

        st_l.update(i, (A[i] + (N - i) * D, i))
        st_r.update(i, (A[i] + i * D, i))

    # return [(i, j, abs(i - j) * D + A[i] + A[j])
    #         for i in range(N) for j in range(i + 1, N)]
    return edges


class SegmentTree:

    # def __init__(self, size: int, op: Callable[[any, any], any], init: any = 0):
    def __init__(self, size, op, init):
        """initialize SegmentTree
        :param size: Size of tree. This must be natural number.
        :param op: Operator which compares two numbers. This function return the
                representative value of two arguments.
        :param init: The initial value of element of tree.
        """
        if size < 1:
            raise Exception('size must be greater than 0 (actual = %d)' % size)

        self.__treesize = 1
        self.__init = init
        # tree size is the minimum number which is greater than or equal to
        # designed `size` and is power of 2.
        while self.__treesize < size:
            self.__treesize *= 2

        self.__size = self.__treesize
        self.__treesize = self.__treesize * 2 - 1

        self.__comp = op

        # initialize tree
        self.__tree = [init for _ in range(self.__treesize)]

    def range(self, l: int, r: int) -> int:
        """Returns the representative value in range [l, r)
        :param l: left value(include)
        :param r: right value(not include)
        :return: the representative value in range[l, r)
        """
        return self.__range(l, r, 0, 0, self.__size)

    def __range(self, l: int, r: int, k: int, kl: int, kr: int) -> int:
        """
        :param l: left value (include)
        :param r: right value (not include)
        :param k: node index
        :param kl: left value of node
        :param kr: right value of node
        """
        if kr <= l or r <= kl:
            # no crossing
            return self.__init
        if l <= kl and kr <= r:
            # including whole k's range
            return self.__tree[k]

        vl = self.__range(l, r, 2 * k + 1, kl, (kl + kr) // 2)
        vr = self.__range(l, r, 2 * k + 2, (kl + kr) // 2, kr)
        if vl is None:
            return vr
        if vr is None:
            return vl

        return self.__comp(vl, vr)

    def update(self, index: int, val: int):
        """update value at self.__tree[index]'s value
        """
        index += self.__size - 1
        self.__tree[index] = val

        while index > 0:
            index = (index - 1) // 2
            self.__tree[index] = self.__comp(
                self.__tree[2*index + 1], self.__tree[2*index + 2])

    def print(self):
        """for debug
        """
        print(self.__tree)


class UnionFindTree:

    def __init__(self, size: int):
        """UnionFindTreeを初期化します
        :param size: サイズ
        :return: UnionFindTree
        """
        if size < 0:
            raise Exception('size must be greater than 0.')

        self.__parent = [-1] * size

    def summarize(self, a: int, b: int):
        """aを含む木とbを含む木をまとめます
        :param a: 木a
        :param b: 木b
        """
        a = self.__root(a)
        b = self.__root(b)

        if a != b:
            self.__parent[b] = a

    def __root(self, a: int) -> int:
        """aの根を求めます
        :param a: 木の要素
        :return: 根
        """
        if self.__parent[a] == -1:
            return a
        else:
            return self.__root(self.__parent[a])

    def same(self, a: int, b: int) -> bool:
        """a と b が同じ木に属するかを判断します
        :param a: 木
        :param b: 木
        :return: a と b が同じ木なら True。そうでないなら False。
        """
        return self.__root(a) == self.__root(b)


if __name__ == "__main__":
    N, D = [int(s) for s in input().split()]
    A = [int(s) for s in input().split()]
    ans = connecting_cities(N, D, A)
    print(ans)

Submission status: Time Limit Exceeded

Solution code:


def connecting_cities(N: int, D: int, A: list) -> int:
    edges = [(i, j, abs(i - j) * D + A[i] + A[j])
             for i in range(N) for j in range(i + 1, N)]
    edges = sorted(edges, key=lambda x: x[2], reverse=True)

    uft = UnionFindTree(N)

    count = 0
    while not len(edges) == 0:
        u, v, d = edges.pop()

        if not uft.same(u, v):
            uft.summarize(u, v)
            count += d

    return count


class UnionFindTree:

    def __init__(self, size: int):
        """UnionFindTreeを初期化します
        :param size: サイズ
        :return: UnionFindTree
        """
        if size < 0:
            raise Exception('size must be greater than 0.')

        self.__parent = [-1] * size

    def summarize(self, a: int, b: int):
        """aを含む木とbを含む木をまとめます
        :param a: 木a
        :param b: 木b
        """
        a = self.__root(a)
        b = self.__root(b)

        if a != b:
            self.__parent[b] = a

    def __root(self, a: int) -> int:
        """aの根を求めます
        :param a: 木の要素
        :return: 根
        """
        if self.__parent[a] == -1:
            return a
        else:
            return self.__root(self.__parent[a])

    def same(self, a: int, b: int) -> bool:
        """a と b が同じ木に属するかを判断します
        :param a: 木
        :param b: 木
        :return: a と b が同じ木なら True。そうでないなら False。
        """
        return self.__root(a) == self.__root(b)


if __name__ == "__main__":
    N, D = [int(s) for s in input().split()]
    A = [int(s) for s in input().split()]
    ans = connecting_cities(N, D, A)
    print(ans)

Submission status: Time Limit Exceeded

Solution code:
# パラメータ入力
N, D = map(int,input().split())
A = list(map(int,input().split()))

# 各都市間のコストを求める
def cost(A, i, j):
    return (abs(i-j) * D) + A[i - 1] + A[j - 1]

cost_list = []
a_sum = len(A)
i = a_sum
j = 1
while i > 0:
    while j < i:
        cost_list.append([cost(A, i, j), i, j])
        j += 1
    j = 1
    i -= 1

# 各都市間のコストを低い順にソートする
cost_list.sort()

# 同じ国リスト間をループする場合は、カウントしない
# 繋いだ国があれば、その国クリストに繋ぐ
# 繋いだ国がなければ、新規に国リストを作る
# 国リスト間を跨ぐ場合は、国リストをマージする
group_map = {}
same_group_map = {}
currnt_group_index = 1
def add_group(country_num_1, country_num_2):
    global group_map
    if (country_num_1 in group_map) == False and (country_num_2 in group_map) == False:
        global currnt_group_index
        group_map[country_num_1] = currnt_group_index
        group_map[country_num_2] = currnt_group_index
        currnt_group_index += 1
    elif (country_num_1 in group_map) == True and (country_num_2 in group_map) == False:
        group_map[country_num_2] = group_map[country_num_1]
    elif (country_num_1 in group_map) == False and (country_num_2 in group_map) == True:
        group_map[country_num_1] = group_map[country_num_2]
    else:
        l = group_map[country_num_1]
        r = group_map[country_num_2]
        if l != r:
            if (l in same_group_map) == False:
                same_group_map[l] = []
            else:
                same_group_map[l].append(r)
            if (r in same_group_map) == False:
                    same_group_map[r] = []
            else:
                same_group_map[l].append(l)

def can_construct(country_num_1, country_num_2):
    if (country_num_1 in group_map) == False:
        return True
    if (country_num_2 in group_map) == False:
        return True

    l = group_map[country_num_1]
    r = group_map[country_num_2]
    if l != r:
        if (l in same_group_map) == False:
            return True
        if (r in same_group_map) == False:
            return True
        if r in same_group_map[l]:
            return False
        if l in same_group_map[r]:
            return False
    return False

# 小さい順からN-1個の分の総和を求める
# その際最低1つ以上の新規国を繋ぐ必要がある
min_cost_sum = 0
count = 0

for road in cost_list:
    contry_1 = road[1]
    contry_2 = road[2]
    if can_construct(contry_1, contry_2) == False:
        continue
    add_group(contry_1, contry_2)
    min_cost_sum += road[0]
    count += 1
    if count == a_sum - 1:
        break

print(min_cost_sum)

Submission status: Time Limit Exceeded

Solution code:
# パラメータ入力
N, D = map(int,input().split())
A = list(map(int,input().split()))

# 各都市間のコストを求める
def cost(A, i, j):
    return (abs(i-j) * D) + A[i - 1] + A[j - 1]

cost_list = []
a_sum = len(A)
i = a_sum
j = 1
while i > 0:
    while j < i:
        cost_list.append([cost(A, i, j), i, j])
        j += 1
    j = 1
    i -= 1

# 各都市間のコストを低い順にソートする
cost_list.sort()

# 同じ国リスト間をループする場合は、カウントしない
# 繋いだ国があれば、その国クリストに繋ぐ
# 繋いだ国がなければ、新規に国リストを作る
# 国リスト間を跨ぐ場合は、国リストをマージする
group_map = {}
currnt_group_index = 1
def add_group(country_num_1, country_num_2):
    global group_map
    if (country_num_1 in group_map) == False and (country_num_2 in group_map) == False:
        global currnt_group_index
        group_map[country_num_1] = currnt_group_index
        group_map[country_num_2] = currnt_group_index
        currnt_group_index += 1
    elif (country_num_1 in group_map) == True and (country_num_2 in group_map) == False:
        group_map[country_num_2] = group_map[country_num_1]
    elif (country_num_1 in group_map) == False and (country_num_2 in group_map) == True:
        group_map[country_num_1] = group_map[country_num_2]
    else:
        target_value = group_map[country_num_2]
        overwrite_value = group_map[country_num_1]
        if group_map[country_num_1] != target_value:
            for key in group_map:
                if group_map[key] == target_value:
                    group_map[key] = overwrite_value

def can_construct(country_num_1, country_num_2):
    if (country_num_1 in group_map) == False:
        return True
    if (country_num_2 in group_map) == False:
        return True
    if group_map[country_num_1] != group_map[country_num_2]:
        return True
    return False

# 小さい順からN-1個の分の総和を求める
# その際最低1つ以上の新規国を繋ぐ必要がある
min_cost_sum = 0
count = 0

for road in cost_list:
    contry_1 = road[1]
    contry_2 = road[2]
    if can_construct(contry_1, contry_2) == False:
        continue
    add_group(contry_1, contry_2)
    min_cost_sum += road[0]
    count += 1
    if count == a_sum - 1:
        break

print(min_cost_sum)

Submission status: Time Limit Exceeded

Solution code:
# パラメータ入力
N, D = map(int,input().split())
A = list(map(int,input().split()))

# 各都市間のコストを求める
def cost(A, i, j):
    return (abs(i-j) * D) + A[i - 1] + A[j - 1]

cost_list = []
a_sum = len(A)
i = a_sum
j = 1
while i > 0:
    while j < i:
        cost_list.append([cost(A, i, j), i, j])
        j += 1
    j = 1
    i -= 1

# 各都市間のコストを低い順にソートする
cost_list.sort()

# 同じ国リスト間をループする場合は、カウントしない
# 繋いだ国があれば、その国クリストに繋ぐ
# 繋いだ国がなければ、新規に国リストを作る
# 国リスト間を跨ぐ場合は、国リストをマージする
group_map = {}
currnt_group_index = 1
def add_group(country_num_1, country_num_2):
    if (country_num_1 in group_map) == False and (country_num_2 in group_map) == False:
        global currnt_group_index
        group_map[country_num_1] = currnt_group_index
        group_map[country_num_2] = currnt_group_index
        currnt_group_index += 1
    elif (country_num_1 in group_map) == True and (country_num_2 in group_map) == False:
        group_map[country_num_2] = group_map[country_num_1]
    elif (country_num_1 in group_map) == False and (country_num_2 in group_map) == True:
        group_map[country_num_1] = group_map[country_num_2]
    else:
        if group_map[country_num_1] != group_map[country_num_2]:
            for key in group_map:
                if group_map[key] == group_map[country_num_2]:
                    group_map[key] = group_map[country_num_1]

def can_construct(country_num_1, country_num_2):
    if (country_num_1 in group_map) == False:
        return True
    if (country_num_2 in group_map) == False:
        return True
    if group_map[country_num_1] != group_map[country_num_2]:
        return True
    return False

# 小さい順からN-1個の分の総和を求める
# その際最低1つ以上の新規国を繋ぐ必要がある
min_cost_sum = 0
count = 0

for road in cost_list:
    contry_1 = road[1]
    contry_2 = road[2]
    if can_construct(contry_1, contry_2) == False:
        continue
    add_group(contry_1, contry_2)
    min_cost_sum += road[0]
    count += 1
    if count == a_sum - 1:
        break

print(min_cost_sum)