from __future__ import annotations

import networkx as nx

from graph.exploration.algorithms.in_memory.graph_features.common import FeatureRowsByKey


def apply_connected_components(rows_by_key: FeatureRowsByKey, graph: nx.Graph) -> None:
    for component_id, component in enumerate(nx.connected_components(graph)):
        component_size = len(component)
        for node_key in component:
            rows_by_key[node_key]["connected_component_id"] = component_id
            rows_by_key[node_key]["connected_component_size"] = component_size
