← run

py-06-numpy-distances

1.000
3/3 tests· math
Challenge · difficulty 3/5
# Pairwise distances (numpy)

Implement **`solution.py`** with:

```python
import numpy as np

def pairwise_distances(points: np.ndarray) -> np.ndarray:
    ...
```

`points` is a 2-D array of shape `(n, d)` (n points in d dimensions). Return an `(n, n)`
array where entry `[i, j]` is the Euclidean distance between point `i` and point `j`.

- The result must be a NumPy array (`np.ndarray`), symmetric, with a zero diagonal.
- Use NumPy vectorization (no Python `for` loops over point pairs).
- Example: for `[[0, 0], [3, 4]]` the result is `[[0, 5], [5, 0]]`.
tests/test_distances.py
import numpy as np
from solution import pairwise_distances


def test_simple():
    pts = np.array([[0.0, 0.0], [3.0, 4.0]])
    d = pairwise_distances(pts)
    assert isinstance(d, np.ndarray)
    assert d.shape == (2, 2)
    np.testing.assert_allclose(d, [[0.0, 5.0], [5.0, 0.0]], atol=1e-9)


def test_symmetric_zero_diagonal():
    rng = np.random.default_rng(0)
    pts = rng.standard_normal((6, 3))
    d = pairwise_distances(pts)
    assert d.shape == (6, 6)
    np.testing.assert_allclose(d, d.T, atol=1e-9)
    np.testing.assert_allclose(np.diag(d), np.zeros(6), atol=1e-9)


def test_matches_bruteforce():
    rng = np.random.default_rng(42)
    pts = rng.standard_normal((5, 4))
    d = pairwise_distances(pts)
    for i in range(5):
        for j in range(5):
            expected = np.sqrt(((pts[i] - pts[j]) ** 2).sum())
            assert abs(d[i, j] - expected) < 1e-9
System prompt
You are an expert programmer. Solve the task exactly as specified. Output your solution as fenced code blocks using the required file name(s) and the exact function/type signatures requested. Prefer correctness; do not include prose outside code unless asked.
Proposed solution
```python
import numpy as np

def pairwise_distances(points: np.ndarray) -> np.ndarray:
    """
    Compute pairwise Euclidean distances between points.

    Args:
        points: A 2-D array of shape (n, d) representing n points in d dimensions.

    Returns:
        An (n, n) array where entry [i, j] is the Euclidean distance between point i and point j.
        The result is symmetric with a zero diagonal.
    """
    # Compute squared distances using broadcasting
    # points[:, None, :] is (n, 1, d)
    # points[None, :, :] is (1, n, d)
    # diff is (n, n, d)
    diff = points[:, None, :] - points[None, :, :]
    
    # Sum over the last axis (dimensions) to get squared Euclidean distances
    sq_dist = np.sum(diff ** 2, axis=-1)
    
    # Take the square root to get Euclidean distances
    return np.sqrt(sq_dist)
```

Test output (stdout)

...                                                                      [100%]
3 passed in 0.05s