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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
| '''
Author: NEFU AB-IN
Date: 2024-07-17 15:46:10
FilePath: \LeetCode\2959\2959.py
LastEditTime: 2024-07-17 16:16:35
'''
import random
from collections import Counter, defaultdict, deque
from datetime import datetime, timedelta
from functools import lru_cache
from heapq import heapify, heappop, heappush, nlargest, nsmallest
from itertools import combinations, compress, permutations, starmap, tee
from math import ceil, comb, fabs, floor, gcd, log, perm, sqrt
from string import ascii_lowercase, ascii_uppercase
from sys import exit, setrecursionlimit, stdin
from typing import Any, Dict, List, Tuple, TypeVar, Union
TYPE = TypeVar('TYPE')
N = int(2e5 + 10)
M = int(20)
INF = int(1e12)
OFFSET = int(100)
MOD = int(1e9 + 7)
setrecursionlimit(int(2e9))
class Arr:
array = staticmethod(lambda x=0, size=N: [x() if callable(x) else x for _ in range(size)])
array2d = staticmethod(lambda x=0, rows=N, cols=M: [Arr.array(x, cols) for _ in range(rows)])
graph = staticmethod(lambda size=N: [[] for _ in range(size)])
class Math:
max = staticmethod(lambda a, b: a if a > b else b)
min = staticmethod(lambda a, b: a if a < b else b)
class IO:
input = staticmethod(lambda: stdin.readline().rstrip("\r\n"))
read = staticmethod(lambda: map(int, IO.input().split()))
read_list = staticmethod(lambda: list(IO.read()))
class Std:
class GraphShortestPath:
def __init__(self, n: int):
self.n = n
self.g = Arr.graph(n)
def add_edge(self, u: int, v: int, w: int):
"""Add an edge to the graph."""
self.g[u].append((v, w))
def floyd(self, exclude_nodes=None) -> List[List[int]]:
"""Floyd's algorithm for finding the shortest paths between all pairs of nodes."""
if exclude_nodes is None:
exclude_nodes = []
dist = Arr.array2d(INF, self.n, self.n)
for u in range(self.n):
if u in exclude_nodes:
continue
for v, w in self.g[u]:
if v in exclude_nodes:
continue
dist[u][v] = Math.min(dist[u][v], w)
for i in range(self.n):
if i in exclude_nodes:
continue
dist[i][i] = 0
for k in range(self.n):
for i in range(self.n):
if dist[i][k] == INF:
continue
for j in range(self.n):
if dist[i][j] > dist[i][k] + dist[k][j]:
dist[i][j] = dist[i][k] + dist[k][j]
return dist
class Solution:
def numberOfSets(self, n: int, maxDistance: int, roads: List[List[int]]) -> int:
def judge(dist, exclude):
for u in range(n):
if u in exclude:
continue
for v in range(n):
if v in exclude:
continue
if dist[u][v] > maxDistance:
return 0
return 1
graph = Std.GraphShortestPath(n)
for u, v, w in roads:
graph.add_edge(u, v, w)
graph.add_edge(v, u, w)
ans = 0
for i in range(1 << n):
exclude = set()
for j in range(n):
if i & 1 << j:
exclude.add(j)
dist = graph.floyd(exclude)
ans += judge(dist, exclude)
return ans
|