#[derive(Debug, Clone, Copy)]
pub struct Edge {
    pub weight: usize,
    pub source_id: usize,
    pub target_id: usize,
}

pub struct WeightedGraph {
    pub n_nodes: usize,
    pub edges: Vec<Edge>,
}

pub fn k_clustering(graph: &mut WeightedGraph) -> usize {
    let k: usize = 4;

    let mut node_id_to_cluster_id: Vec<usize> = (0..graph.n_nodes).collect();
    let mut cluster_count = graph.n_nodes;
    let mut clusters: Vec<Vec<usize>> = (0..graph.n_nodes).map(|x| vec![x]).collect();

    // Sort edges by increasing weight
    graph
        .edges
        .sort_by(|a, b| a.weight.partial_cmp(&b.weight).unwrap());

    for edge in &graph.edges {
        let cluster_id_a = node_id_to_cluster_id[edge.source_id];
        let cluster_id_b = node_id_to_cluster_id[edge.target_id];

        if cluster_id_a != cluster_id_b {
	        let mut cluster_b = std::mem::take(&mut clusters[cluster_id_b]);
	        for node_id in &cluster_b {
	        	node_id_to_cluster_id[*node_id] = cluster_id_a;
	        }
	        clusters[cluster_id_a].append(&mut cluster_b);
	        cluster_count -= 1;
        }

        if cluster_count <= k {
            break;
        }
    }

    let mut min_spacing = usize::MAX;
    for edge in &graph.edges {
        let cluster_id_a = node_id_to_cluster_id[edge.source_id];
        let cluster_id_b = node_id_to_cluster_id[edge.target_id];
        if cluster_id_a != cluster_id_b {
        	if edge.weight < min_spacing {
        		min_spacing = edge.weight;
        	}
        }
    }

    min_spacing
}