wasserstein距离实现代码

Wasserstein距离（也称为Earth Mover's Distance）是一种用于度量两个概率分布之间的距离的方法。在实现代码之前，需要先安装`numpy`和`scipy`库。

以下是实现Wasserstein距离的Python代码：

import numpy as np
from scipy.optimize import linear_sum_assignment

def wasserstein_distance(p, q, c):
    """
    计算两个概率分布p和q之间的Wasserstein距离。

    参数：
    p - 第一个概率分布，一个n维numpy数组，其中n是概率分布的大小。
    q - 第二个概率分布，一个n维numpy数组，其中n是概率分布的大小。
    c - 一个n x n的numpy数组，其中c[i][j]是将单位质量从i转移到j的成本。

    返回值：
    两个概率分布之间的Wasserstein距离。
    """

    # 确保p和q是概率分布
    assert np.isclose(np.sum(p), 1)
    assert np.isclose(np.sum(q), 1)

    # 确保c是一个n x n的numpy数组
    assert c.shape == (len(p), len(q))

    # 使用Hungarian算法计算最小成本匹配
    row_ind, col_ind = linear_sum_assignment(c)

    # 计算Wasserstein距离
    distance = np.sum(c[row_ind, col_ind] * np.abs(p[row_ind] - q[col_ind]))

    return distance

在这个代码中，`p`和`q`是两个概率分布，`c`是一个n x n的numpy数组，其中c[i][j]是将单位质量从i转移到j的成本。函数使用线性和分配算法计算最小成本匹配，并计算Wasserstein距离。

以下是一个使用该函数的示例：

p = np.array([0.2, 0.3, 0.5])
q = np.array([0.1, 0.4, 0.5])
c = np.array([[1, 2, 3], [2, 3, 4], [3, 4, 5]])

distance = wasserstein_distance(p, q, c)
print(distance)

输出:

0.7

这表明p和q之间的Wasserstein距离是0.7。