Created
April 3, 2024 19:54
-
-
Save coderodde/159b89d4484342b1c69f852599ff648f to your computer and use it in GitHub Desktop.
Colorful subgraph problem.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from itertools import combinations | |
import time | |
def millis(): | |
return round(1000 * time.time()) | |
def powerset(iterable): | |
s = list(iterable) | |
list_all_combinations = [] | |
for r in range(len(s) + 1): | |
list_all_combinations += list(combinations(s, r)) | |
return list_all_combinations | |
def dfs(v, visited, subgraph_set, g, k): | |
if len(visited) > k: | |
return | |
visited.add(v) | |
for u in g[v]: | |
if u not in visited: | |
dfs(u, visited, subgraph_set, g, k) | |
def create_graph(graph_info): | |
vertices: list[str] | |
edges: list[tuple[str, str]] | |
vertices, edges = graph_info['vertices'], graph_info['edges'] | |
v_to_idx: dict[str, int] = {v: idx for idx, v in enumerate(vertices)} # idx_to_v is not needed, it's vertices | |
color_map: dict[str, int] = graph_info['color_constraints'] | |
num_colors: int = graph_info['num_colors'] | |
# Create (undirected) graph: | |
g = {v: set() for v in vertices} | |
for u, v in edges: | |
g[u].add(v) | |
g[v].add(u) | |
return vertices, edges, v_to_idx, color_map, num_colors, g | |
def find_colorful_subtree(graph_info) -> bool: | |
vertices, edges, v_to_idx, color_map, num_colors, g = create_graph(graph_info) | |
# DP table dp[v][S] = True if a subtree rooted at v is S-colorful | |
# We encode S with integers (bitmask) | |
dp = [[False for _ in range(1 << num_colors)] for _ in range(len(vertices))] | |
# Initialize the DP table | |
for v_idx, v in enumerate(vertices): | |
color = color_map[v] | |
dp[v_idx][1 << color - 1] = True | |
# Recursive function to fill the DP table | |
for subset in powerset(range(num_colors)): | |
if len(subset) <= 1: | |
continue | |
# generate all s1, s2 such that s1 + s2 = subset and s1 & s2 = {} | |
s: int = 0 | |
for color in subset: | |
s |= 1 << color | |
for left_subset in powerset(subset): | |
if len(left_subset) == 0 or len(left_subset) == len(subset): | |
continue | |
left_mask: int = 0 | |
for color in left_subset: | |
left_mask |= 1 << color | |
right_mask: int = s ^ left_mask # ^ -> XOR | |
for v_idx, v in enumerate(vertices): | |
if (1 << color_map[v] - 1) & left_mask != 0: | |
for u in g[v]: | |
if (1 << color_map[u] - 1) & right_mask == 0: | |
continue | |
if dp[v_idx][left_mask] and dp[v_to_idx[u]][right_mask]: | |
dp[v_idx][s] = True | |
break # no need to check other neighbors, we found a colorful subtree rooted at v | |
ans: bool = False | |
for v_idx, v in enumerate(vertices): | |
ans |= dp[v_idx][(1 << num_colors) - 1] | |
return ans | |
def find_colorful_subtree_naive(graph_info) -> bool: | |
vertices, edges, v_to_idx, color_map, num_colors, g = create_graph(graph_info) | |
# we take a set of vertices (2^(|V|)) and check if it is colorful (O(|V|)) | |
for s in range(1 << len(vertices)): | |
if bin(s).count('1') < num_colors: | |
continue | |
# Check if the set is colorful | |
colors: set[int] = set() | |
for v_idx, v in enumerate(vertices): | |
if (1 << v_idx) & s: | |
if color_map[v] in colors: | |
break | |
colors.add(color_map[v]) | |
if len(colors) == num_colors: | |
# Check if the set of vertices is connected | |
visited: set[str] = set() | |
def dfs(v: str): | |
visited.add(v) | |
for u in g[v]: | |
if u not in visited and (1 << v_to_idx[u]) & s: | |
dfs(u) | |
for v_idx, v in enumerate(vertices): | |
if (1 << v_idx) & s: | |
dfs(v) | |
break | |
if len(visited) == bin(s).count('1'): | |
return True | |
return False | |
example_input_1 = { | |
"vertices": ["v1", "v2", "v3", "v4", "v5"], | |
"edges": [ | |
("v1", "v2"), | |
("v2", "v3"), | |
("v3", "v4"), | |
("v4", "v5"), | |
("v1", "v3"), | |
("v2", "v4"), | |
], | |
"color_constraints": {"v1": 1, "v2": 2, "v3": 3, "v4": 1, "v5": 2}, | |
"num_colors": 3, | |
} # Should be True | |
example_input_2 = { | |
"vertices": ["v1", "v2", "v3", "v4", "v5"], | |
"edges": [ | |
("v1", "v2"), | |
("v2", "v3"), | |
("v3", "v4"), | |
("v4", "v5"), | |
("v1", "v3"), | |
("v2", "v4"), | |
], | |
"color_constraints": {"v1": 2, "v2": 1, "v3": 1, "v4": 1, "v5": 3}, | |
"num_colors": 3, | |
} # Should be False | |
example_input_3 = { | |
"vertices": ["v1", "v2", "v3", "v4", "v5"], | |
"edges": [ | |
("v1", "v2"), | |
("v2", "v3"), | |
("v3", "v4"), | |
("v4", "v5"), | |
("v1", "v3"), | |
("v2", "v4"), | |
], | |
"color_constraints": {"v1": 2, "v2": 1, "v3": 1, "v4": 2, "v5": 3}, | |
"num_colors": 3, | |
} # should be True | |
class IndexIterator: | |
def __init__(self, n, k): | |
self.n = n | |
self.k = k | |
self.indices = [] | |
self.index_set = set() | |
for i in range(k): | |
self.indices.append(i) | |
self.index_set.add(i) | |
def __str__(self): | |
return ','.join(str(x) for x in self.indices) | |
def get_indices(self): | |
return self.indices | |
def get_index_set(self): | |
return self.index_set | |
def increment(self): | |
if self.indices[self.k - 1] < self.n - 1: | |
self.index_set.remove(self.indices[self.k - 1]) | |
self.index_set.add(self.indices[self.k - 1] + 1) | |
self.indices[self.k - 1] += 1 | |
return True | |
for i in range(self.k - 2, -1, -1): | |
if self.indices[i] < self.indices[i + 1] - 1: | |
self.index_set.remove(self.indices[i]) | |
self.index_set.add(self.indices[i] + 1) | |
self.indices[i] += 1 | |
for j in range(i + 1, self.k): | |
self.index_set.remove(self.indices[j]) | |
self.index_set.add(self.indices[j - 1] + 1) | |
self.indices[j] = self.indices[j - 1] + 1 | |
return True | |
return False | |
def compute_graph_inverse(vertices): | |
g = {} | |
for idx in range(len(vertices)): | |
g[idx] = vertices[idx]; | |
return g | |
def is_colorful(it, k, g, inverse_g): | |
visited = set() | |
subgraph_node_indices = it.get_indices() | |
subgraph_node_set = set() | |
initial_node = None | |
for idx in subgraph_node_indices: | |
subgraph_node_set.add(inverse_g[idx]) | |
if not initial_node: | |
initial_node = inverse_g[idx] | |
dfs(initial_node, visited, subgraph_node_set, g, k) | |
return len(visited) == k | |
def coderodde_colorful_subgraph(graph_info): | |
vertices, edges, v_to_idx, color_map, num_colors, g = create_graph(graph_info) | |
it = IndexIterator(len(vertices), num_colors) | |
inverse_g = compute_graph_inverse(vertices) | |
while True: | |
if is_colorful(it, num_colors, g, inverse_g): | |
return True | |
if not it.increment(): | |
return False | |
assert coderodde_colorful_subgraph(example_input_1) | |
assert not coderodde_colorful_subgraph(example_input_2) | |
assert coderodde_colorful_subgraph(example_input_3) | |
start = millis() | |
for _ in range(10000): | |
assert find_colorful_subtree(example_input_1) | |
assert not find_colorful_subtree(example_input_2) | |
assert find_colorful_subtree(example_input_3) | |
end = millis() | |
print("find_colorful_subtree in %d milliseconds." % (end - start)) | |
start = millis() | |
for _ in range(10000): | |
assert find_colorful_subtree_naive(example_input_1) | |
assert not find_colorful_subtree_naive(example_input_2) | |
assert find_colorful_subtree_naive(example_input_3) | |
end = millis() | |
print("find_colorful_subtree_naive in %d milliseconds." % (end - start)) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment