Untitled

from collections import defaultdict
from copy import deepcopy
from typing import List, Dict

from test_cases import test_cases


class Solution:

    def dfs(self, nodes, v, seen: set):
        seen.add(v)
        for u in nodes[v]:
            if u[0] not in seen:
                self.dfs(nodes, u[0], seen)

    def get_count_of_connectivity_component(self, graph):
        seen = set()
        comps = 0

        for v in deepcopy(graph):
            if v not in seen:
                self.dfs(graph, v, seen)
                comps += 1
        return comps

    def remove_all_roads_less_than_x(self, graph, x) -> Dict:
        for start, roads in graph.items():
            graph[start] = list(filter(lambda i: i[1] > x, roads))
        return graph

    def find_x(self, n, m, roads: List[List[int]]) -> int:
        graph = defaultdict(list)
        for start, end, weight in roads:
            graph[start].append((end, weight))
            graph[end].append((start, weight))

        base_count_of_states = self.get_count_of_connectivity_component(graph)
        sorted_roads = sorted(roads, key=lambda x1: x1[2])

        left, right = 0, n - 1

        x = -1
        while left <= right:
            midd = (right + left) // 2
            tmp_x: int = sorted_roads[midd][2]

            tmp_graph = self.remove_all_roads_less_than_x(deepcopy(graph), tmp_x)
            new_count_of_states = self.get_count_of_connectivity_component(tmp_graph)

            if new_count_of_states > base_count_of_states:
                right = midd - 1
                x = tmp_x - 1
            else:
                left = midd + 1

        return x


test_cases(
    func=Solution().find_x,
    keyses=['n', 'm', 'roads'],
    params=[
        (2, 2, [[1, 2, 6], [2, 1, 9]]),
        (2, 2, [[1, 2, 6], [2, 1, 9], [3, 4, 9]]),
        (5, 6, [[1, 2, 8], [2, 3, 6], [2, 3, 2], [3, 1, 4], [5, 4, 1], [4, 5, 8]]),
    ],
    answers=[8, 8, 5]
)