# インストール、インポート
!pip install dwave-ocean-sdk
!pip install pyqubo

import numpy as np
import dimod

from dwave.samplers import SimulatedAnnealingSampler

# 都市数
N = 5

# 都市間の移動距離
Q = np.array([[1000, 20, 20, 50, 40], \
              [30, 1000, 10, 30, 20], \
              [20, 10, 1000, 30, 20], \
              [50, 30, 20, 1000, 10], \
              [40, 20, 20, 10, 1000]])

# 変数を作成
from pyqubo import Array, Constraint, Placeholder
x = Array.create('x', shape=(N, N), vartype='BINARY')

# コスト関数
cost = 0
for t in range(N):
  for i in range(N):
    for j in range(N):
      cost += Q[i][j]*x[t][i]*x[(t+1)%N][j]

# 制約条件1：各都市に1回は訪問しなければならない
constr_1 = 0
for i in range(N):
  constr_1 += (np.sum(x[i])-1)**2

# 制約条件2：1度に訪れる都市は1つでなければならない
constr_2 = 0
for i in range(N):
  constr_2 += (np.sum(x.T[i])-1)**2

cost_func = cost + Placeholder('lam')*Constraint(constr_1, label='constr_1') \
                 + Placeholder('lam')*Constraint(constr_2, label='constr_2')
model = cost_func.compile()

feed_dict = {'lam': 1000} # penalty係数
qubo, offset = model.to_qubo(feed_dict=feed_dict)

sampler = SimulatedAnnealingSampler()
sampleset = sampler.sample_qubo(qubo, num_reads=10)

# 結果出力
print(sampleset.record)

# 整形表示

for i in range (10) :
    print("\n",sampleset.record[i][0].reshape(N, N))