random.py

from .utils import get_func
from .degree import node_degree
from .permute import permute


def random_cluster(edge_index, node_rid=None, edge_rid=None, num_nodes=None):
    num_nodes = edge_index.max() + 1 if num_nodes is None else num_nodes
    row, col = permute(edge_index, num_nodes, node_rid, edge_rid)
    degree = node_degree(row, num_nodes, out=row.new())

    cluster = edge_index.new(num_nodes).fill_(-1)
    func = get_func('random', cluster)
    func(cluster, row, col, degree)

    return cluster