import torch


def partial_distance_matrix(x: torch.Tensor, partial_idx: torch.Tensor) -> torch.Tensor:
    """Calculates the partial distance matrix given the indices. For a low memory footprint (small computation graph)
    it is essential to avoid duplicated computation of the distances.

    Parameters
    ----------
    x : torch.Tensor
        Dense [n, d] tensor with attributes to calculate the distance between.
    partial_idx : torch.Tensor
        Dense [batch_size, k] tensor where `-1` stands for no index.
        Pairs are generated by the row id and the contained ids.

    Returns
    -------
    torch.Tensor
        [n, k, k] distances matrix (zero entries for `-1` indices)
    """
    n, _ = x.shape
    batch_size, k = partial_idx.shape

    # Permute the indices of partial_idx
    idx_row = partial_idx[:, None, :].expand(batch_size, k, k).flatten()
    idx_column = partial_idx[:, None, :].expand(batch_size, k, k).transpose(1, 2).flatten()
    is_not_missing_mask = (idx_row != -1) & (idx_column != -1)
    idx_row, idx_column = idx_row[is_not_missing_mask], idx_column[is_not_missing_mask]

    # Use symmetry of Euclidean distance to half memory footprint
    symmetry_mask = idx_column < idx_row
    idx_row[symmetry_mask], idx_column[symmetry_mask] = idx_column[symmetry_mask], idx_row[symmetry_mask]
    del symmetry_mask

    # Create linear index (faster deduplication)
    linear_index = idx_row * n + idx_column
    del idx_row
    del idx_column

    # Avoid duplicated distance calculation (helps greatly for space cost of backward)
    distance_matrix_idx, unique_reverse_index = torch.unique(linear_index, return_inverse=True)

    # Calculate Euclidean distances between all pairs
    sparse_distances = torch.norm(x[distance_matrix_idx // n] - x[distance_matrix_idx % n], dim=1)

    # Create dense output
    out = torch.zeros(batch_size * k * k, dtype=torch.float, device=x.device)

    # Map sparse distances to output tensor
    out[is_not_missing_mask] = sparse_distances[unique_reverse_index]

    return out.view(batch_size, k, k)