Day 025-Q5 — 最適輸送 (EMD / Wasserstein) + Sinkhorn

問題

2 つの離散確率分布 $\mu, \nu$ について Earth Mover's Distance ($W_1$ Wasserstein) を計算。さらに Sinkhorn 反復によるエントロピー正則化近似解も求める。

制約

入出力例

2 2
0 0 1 2
2 0 1 2
1 0 1 2
3 0 1 2
0.1 100

Exact EMD (min-cost flow): 1.0000
Sinkhorn EMD (eps=0.1, iter=100): X.XXXX
Convergence achieved at iteration: XX

ヒント (段階的開示)

模範解答 (Python)

📋 コピーimport sys
import numpy as np
from scipy.optimize import linprog
input = sys.stdin.readline

def main():
    N, M = map(int, input().split())
    mu_points = []; mu_weights = []
    for _ in range(N):
        line = input().split()
        x, y = int(line[0]), int(line[1])
        w = int(line[2]) / int(line[3])
        mu_points.append((x, y)); mu_weights.append(w)
    nu_points = []; nu_weights = []
    for _ in range(M):
        line = input().split()
        a, b = int(line[0]), int(line[1])
        v = int(line[2]) / int(line[3])
        nu_points.append((a, b)); nu_weights.append(v)
    eps_str, max_iter_str = input().split()
    eps = float(eps_str); max_iter = int(max_iter_str)

    mu = np.array(mu_weights, dtype=float)
    nu = np.array(nu_weights, dtype=float)
    C = np.zeros((N, M))
    for i in range(N):
        for j in range(M):
            dx = mu_points[i][0] - nu_points[j][0]
            dy = mu_points[i][1] - nu_points[j][1]
            C[i][j] = (dx*dx + dy*dy) ** 0.5

    c_flat = C.flatten()
    A_eq_rows = []; b_eq = []
    for i in range(N):
        row = np.zeros(N * M)
        for j in range(M):
            row[i * M + j] = 1
        A_eq_rows.append(row); b_eq.append(mu[i])
    for j in range(M):
        row = np.zeros(N * M)
        for i in range(N):
            row[i * M + j] = 1
        A_eq_rows.append(row); b_eq.append(nu[j])
    A_eq = np.array(A_eq_rows); b_eq = np.array(b_eq)
    bounds = [(0, None)] * (N * M)
    result = linprog(c_flat, A_eq=A_eq, b_eq=b_eq, bounds=bounds, method='highs')
    exact_emd = result.fun if result.success else float('nan')
    print(f"Exact EMD (min-cost flow): {exact_emd:.4f}")

    K = np.exp(-C / eps)
    u = np.ones(N); v_vec = np.ones(M)
    converged_iter = max_iter; tol = 1e-9
    for iteration in range(max_iter):
        u_new = mu / (K @ v_vec + 1e-300)
        v_new = nu / (K.T @ u_new + 1e-300)
        diff = np.max(np.abs(u_new - u))
        u = u_new; v_vec = v_new
        if diff < tol:
            converged_iter = iteration + 1; break
    gamma = np.diag(u) @ K @ np.diag(v_vec)
    sinkhorn_cost = np.sum(gamma * C)
    print(f"Sinkhorn EMD (eps={eps}, iter={max_iter}): {sinkhorn_cost:.4f}")
    print(f"Convergence achieved at iteration: {converged_iter}")

main()

Step-by-Step 解説

よくあるミス

次のステップ

• $W_2$ Wasserstein と半離散 OT
• Wasserstein Barycenter
• Kantorovich-Rubinstein 双対

自己評価

自分の回答

気づき・メモ