如何从Scipy的wasserstein_distance中提取距离和传输矩阵？

import numpy as np
from scipy.stats import wasserstein_distance
from scipy.optimize import linear_sum_assignment

# 假设有两个概率分布P和Q
P = np.array([0.2, 0.3, 0.5])  # 第一个概率分布P
Q = np.array([0.1, 0.4, 0.5])  # 第二个概率分布Q

distance = wasserstein_distance(P, Q)
print("Wasserstein距离：", distance)

# 构建距离矩阵
distance_matrix = np.zeros((len(P), len(Q)))
for i in range(len(P)):
    for j in range(len(Q)):
        distance_matrix[i, j] = np.abs(i - j)

# 使用linear_sum_assignment函数进行最优分配
row_ind, col_ind = linear_sum_assignment(distance_matrix)
transport_matrix = np.zeros((len(P), len(Q)))
transport_matrix[row_ind, col_ind] = 1

print("传输矩阵：")
print(transport_matrix)