import torch
from torch_geometric.utils import normalized_cut


def test_normalized_cut():
    row = torch.LongTensor([0, 1, 1, 1, 2, 2, 3, 3, 4, 4])
    col = torch.LongTensor([1, 0, 2, 3, 1, 4, 1, 4, 2, 3])
    edge_attr = torch.Tensor([3, 3, 6, 3, 6, 1, 3, 2, 1, 2])
    expected_output = [4, 4, 5, 2.5, 5, 1, 2.5, 2, 1, 2]

    output = normalized_cut(torch.stack([row, col], dim=0), edge_attr)
    assert output.tolist() == expected_output
