Computing All Optimal Partial $p$-Wasserstein Matchings on the Line

Abstract: For $p \ge 1$, the $p$-Wasserstein distance measures the minimum cost of transporting probability mass between distributions, where moving unit mass between two points costs the $p$th power of their distance. For discrete distributions in one dimension, full transport is especially simple: after sorting, mass is matched in order along the line. By contrast, partial and unbalanced transport on the line remains much less understood. Recently, Chapel and Tavenard [ICLR'25] showed that, for $p=1$, all optimal partial transport plans between distributions supported on $n$ points, with uniform mass at each point, can be computed in $O(n\log n)$ time by exploiting the metric structure of the cost. For $p&gt;1$, this structure no longer applies, and existing approaches require $Ω(n<sup>2)$ time. Our main contribution is an FFT-based data structure for balanced-interval transport queries, which bypasses this quadratic bottleneck and yields an $O(p\,n\log² n)$-time algorithm for computing all optimal partial transports on the line for every finite $p\ge 1$. We also provide an open-source C++ implementation that outperforms the state-of-the-art baseline on a range of synthetic instances. Finally, we establish a conditional lower bound for $p=\infty$: any subquadratic-time algorithm for computing all optimal partial transport plan costs on the line would violate the $(\min,+)$-Convolution Hypothesis. This separates the problem from full optimal transport, which is solvable in $O(n\log n)$.

• The paper presents an O(p n log² n)-time, O(n log n)-space algorithm for computing every partial p-Wasserstein cost on the line when p is a finite integer, improving the previous quadratic-time bound.
• The method combines compact balanced-interval representations with an FFT-preprocessed range tree that answers interval-matching cost queries in O(log n) time, avoiding explicit storage of the Θ(n²) matching edges.
• The paper proves a conditional near-quadratic barrier for p=∞ and stronger hardness results for two-interval queries, while experiments show orders-of-magnitude speedups over prior methods on large synthetic datasets.

Problem and context

The $p$-Wasserstein distance between two probability measures is the $p$th root of the minimum cost of transporting one measure into the other, where moving mass $\delta$ between points at distance $d$ costs $\delta\, d^p$. For discrete distributions supported on $n$ points with uniform mass, the optimal transport profile (OT-profile) maps each $\alpha \in [0,1]$ to the minimum cost of a plan transporting mass $\alpha$; equivalently, it is the sequence $C_1,\dots,C_n$, where $C_k$ is the minimum cost of any $p$0-matching between the two point sets. OT-profiles are used in noise filtering, positive-unlabeled learning, and the robust partial Wasserstein metric.

Full optimal transport on the line is trivial: after sorting, matching in order along the line is optimal for all $p$1. The partial case is not. Chapel and Tavenard [ICLR'25] showed that for $p$2 the entire OT-profile on the line can be computed in $p$3 time; their approach exploits the metric structure of the cost. For every $p$4, however, their method requires $p$5 time. The quadratic barrier has a structural cause: for $p$6, consecutive optimal partial matchings share edges (only $p$7 distinct edges appear across all $p$8), whereas for $p$9 convexity makes each optimal $\delta$0-matching unique, so that $\delta$1 and $\delta$2 share no edges and the sequence of matchings contains $\delta$3 distinct edges. Any algorithm that explicitly materializes all partial matchings therefore incurs quadratic work.

Main result

The paper resolves this question for all finite integer $\delta$4: given point sets $\delta$5 of size $\delta$6, the $\delta$7-partial $\delta$8-Wasserstein distances for all $\delta$9 are computable in $d$0 time using $d$1 space. This matches, up to a logarithmic factor, the complexity of full one-dimensional transport, and improves the prior state of the art from $d$2 to near-linear for every finite $d$3.

The algorithm follows the successive-shortest-path structure of the Hungarian algorithm specialized to the line, as in Chapel–Tavenard, but replaces the quadratic interval-cost computation with a new data structure. Two ingredients matter:

• Compact interval representation. A minimum-cost $d$4-matching is represented by a set of at most $d$5 disjoint balanced intervals (intervals containing equal numbers of red and blue points). Although the union of all partial matchings contains $d$6 edges, this representation uses only $d$7 space, and the minimum-cost $d$8-matching differs from the $d$9-matching by modifying only a constant number of intervals.
• Balanced-interval queries. Identifying which intervals to modify reduces to evaluating the minimum matching cost inside balanced intervals. The paper's main technical contribution is a data structure built in $\delta\, d^p$0 time and $\delta\, d^p$1 space that answers such queries in $\delta\, d^p$2 time.

The net-cost of an augmenting step then takes a simple closed form: for an adjacent free pair $\delta\, d^p$3, it equals $\delta\, d^p$4 if no stored interval lies inside $\delta\, d^p$5, and otherwise $\delta\, d^p$6 for the unique stored interval $\delta\, d^p$7 inside $\delta\, d^p$8. Each iteration performs a constant number of heap operations and $\delta\, d^p$9 queries, so the $n$0 iterations take $n$1 total; preprocessing dominates.

The FFT-based range tree

Inside any balanced interval, the optimal matching pairs points in order, so if the leftmost red point has index $n$2 and the leftmost blue point index $n$3, every matched pair has shift $n$4. The data structure is a binary range tree over the merged sorted order of all $n$5 points. Each node $n$6 stores, for every relevant shift $n$7, three quantities: the total cost of shift-$n$8 pairs crossing from its left child to its right child, symmetrically from right to left, and the total cost $n$9 of all shift-$\alpha \in [0,1]$0 pairs contained in $\alpha \in [0,1]$1. Since relevant shifts number $\alpha \in [0,1]$2 per node, total space is $\alpha \in [0,1]$3.

A query determines the shift $\alpha \in [0,1]$4 of the input interval from its leftmost red and blue points, then aggregates precomputed values over the $\alpha \in [0,1]$5 maximal nodes covering the interval. Correctness rests on a non-obvious decomposability property: when a node's children both intersect the query interval, every shift-$\alpha \in [0,1]$6 pair crossing the node's midline lies entirely within the query interval — a consequence of balancedness combined with the order structure. Matching costs are not decomposable in general, so this property is what makes the aggregation valid.

Preprocessing avoids the naive $\alpha \in [0,1]$7 enumeration of shift values via FFT: since all red points in a left child lie left of all blue points in the right child, cross-term costs expand through the binomial identity $\alpha \in [0,1]$8, and each family of sums is a coefficient of a polynomial product computable by convolution. One FFT per binomial term per node yields $\alpha \in [0,1]$9 per node and $\alpha$0 overall; the $\alpha$1 values follow bottom-up in $\alpha$2.

A practical caveat: the running time carries a factor $\alpha$3, and the construction relies on expanding $\alpha$4, so the approach is specific to integer $\alpha$5; nothing in the upper bound extends to non-integer or infinite $\alpha$6.

Implementation and experiments

The authors provide an open-source C++ implementation (FFTHungarian) that computes interval costs naively below a threshold of 2048 points and queries the data structure above it. Against PAWL, the optimized numba-based implementation of Chapel–Tavenard, on synthetic instances ($\alpha$7) spanning normal distributions with differing means or standard deviations and unimodal-versus-bimodal mixtures, FFTHungarian is never substantially slower and becomes orders of magnitude faster once $\alpha$8 and the distributions differ. Notably, even a 10% difference in mean or standard deviation drives PAWL toward its quadratic worst case while FFTHungarian scales near-linearly. The experimental regime deliberately favors the baseline (substantial distributional overlap), strengthening the comparison; the evaluation covers synthetic data only, and hardware was a single laptop-class machine.

Conditional lower bound for $\alpha$9

For $C_1,\dots,C_n$0, the picture changes sharply. Assuming the $C_1,\dots,C_n$1-Convolution Conjecture, no algorithm computes all $C_1,\dots,C_n$2-partial $C_1,\dots,C_n$3-Wasserstein matching costs for $C_1,\dots,C_n$4 integer points in $C_1,\dots,C_n$5 in $C_1,\dots,C_n$6 time for any $C_1,\dots,C_n$7. The reduction first maps arbitrary $C_1,\dots,C_n$8-convolution instances to monotone $C_1,\dots,C_n$9-convolution instances via the transformation $C_k$0, $C_k$1, then observes that for increasing $C_k$2 and decreasing $C_k$3 with $C_k$4, the $C_k$5-partial $C_k$6-Wasserstein cost equals exactly the $C_k$7-convolution entry $C_k$8.

This yields a clean separation: full $C_k$9-Wasserstein transport admits subquadratic exact algorithms even in two dimensions, whereas computing the entire partial profile is conditionally quadratic-hard already on the line. The lower bound also indicates that the dependence on $p$00 cannot be removed by refining the upper-bound techniques, since the same framework applied at $p$01 would require $p$02-convolutions.

Two further lower bounds address a strictly more general data-structure problem, the Two Intervals Matching Problem, where queries specify separate intervals for each color. Via a gadget reduction from the Hamming Distance Oracle problem (encoding bits as four-point gadgets whose Wasserstein distances distinguish equal from unequal gadgets), the paper shows: (i) unless the Combinatorial Matrix Multiplication Conjecture fails, no combinatorial data structure achieves $p$03 for any $p$04; and (ii) unconditionally up to the matrix multiplication exponent, $p$05. The contrast with the single-interval setting — where near-linear preprocessing and $p$06 queries are achieved — shows that allowing two independent query intervals provably increases the problem's difficulty.

Limitations and open questions

Several restrictions bound the applicability of the results. The upper bound requires $p$07 to be a positive integer, since both the running-time dependence and the binomial-expansion technique presuppose it; whether near-linear algorithms exist for fractional $p$08 remains open. The lower bound for $p$09 is conditional on the $p$10-Convolution Conjecture, and the Two Intervals bounds rest on matrix-multiplication conjectures, so unconditional hardness is not established anywhere in the paper. The experimental validation is confined to synthetic one-dimensional distributions with substantial overlap; behavior on real data or on distributions with little overlap is not reported. Finally, the paper leaves open whether the extra $p$11 factor can be removed — the authors note that an attempted unpacking of the FFT did not yield an improvement — and whether the imbalanced (unbalanced-mass) variant admits comparable guarantees; a sketch of a persistent-tree approach for imbalanced single-interval queries appears only in preliminary form without complete proofs.

Conclusion

This paper establishes that the entire OT-profile under the $p$12-Wasserstein cost on the line is computable in $p$13 time for every finite integer $p$14, removing the quadratic barrier that previously applied for all $p$15. The key enabler is an FFT-preprocessed range tree supporting balanced-interval matching queries in logarithmic time, exploiting a decomposability property special to balanced one-dimensional intervals. Together with the conditional quadratic hardness at $p$16 and the stronger hardness for two-interval queries, the results delineate precisely where partial optimal transport on the line retains the computational simplicity of full transport and where it does not.