learning/l33tcode/closest-node-to-path-in-tree.py at master · fspv/learning · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
from typing import List, Set


class Solution:
    def closestNode(
        self, n: int, edges: List[List[int]], query: List[List[int]]
    ) -> List[int]:
        def gen_adj_list(size: int, edges: List[List[int]]) -> List[Set[int]]:
            adj_list: List[Set[int]] = [set() for _ in range(size)]

            for src, dst in edges:
                adj_list[src].add(dst)
                adj_list[dst].add(src)

            return adj_list

        def make_directed_tree(adj_list: List[Set[int]]) -> None:
            def _dfs(node: int) -> None:
                for child in adj_list[node]:
                    adj_list[child].discard(node)
                    _dfs(child)

            _dfs(0)

        adj_list: List[Set[int]] = gen_adj_list(n, edges)
        make_directed_tree(adj_list)

        def lca(first: int, second: int) -> int:
            result = -1

            def _dfs(node: int) -> int:
                nonlocal result

                count = 0

                if node == first:
                    count += 1

                if node == second:
                    count += 1

                for child in adj_list[node]:
                    if count == 2:
                        break

                    count += _dfs(child)

                    if count == 2:
                        break

                if count == 2 and result == -1:
                    result = node

                return count

            _dfs(0)

            return result

        result: List[int] = []

        for src, dst, node in query:
            lca_src_dst = lca(src, dst)
            lca_src_node = lca(src, node)
            lca_dst_node = lca(dst, node)

            """
            sd-sn sd-dn sn-dn
             =      =     =   root
             =      =    !=   invalid
             =     !=    !=   dn
             !=    !=    !=   invalid
             !=    !=     =   sn or dn
             !=     =     =   invalid
             !=     =    !=   sn
              =    !=     =   invalid
            """

            if lca_src_dst == lca_src_node == lca_dst_node:
                result.append(lca_src_dst)
            elif (
                lca_src_dst != lca_src_node
                and lca_src_dst == lca_dst_node
                and lca_src_node != lca_dst_node
            ):
                result.append(lca_src_node)
            elif (
                lca_src_dst != lca_dst_node
                and lca_src_dst == lca_src_node
                and lca_src_node != lca_dst_node
            ):
                result.append(lca_dst_node)
            elif (
                lca_src_dst != lca_src_node
                and lca_src_dst != lca_dst_node
                and lca_src_node == lca_dst_node
            ):
                result.append(lca_src_dst)

        return result