Colorful subgraph problem.

colorful_subgraph.py

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))