twosat.py

from teflib import graph

# N TwoSat
# I {"version": "1.3", "abc": ["Iterable"], "func": ["graph.strongly_connected_component"], "import": ["itertools"]}
class TwoSat:
    '''2-SAT solver.'''

    __slot__ = ('_n', '_graph', '_ord', '_nord')

    def __init__(self, var_count: int) -> None:
        self._n = var_count
        self._graph = [set() for _ in range(var_count * 2)]
        self._ord = self._nord = None

    def x_is_true(self, x: int) -> None:
        x = x + x if x >= 0 else -1 - x - x
        self._graph[x ^ 1].add(x)

    def x_or_y(self, x: int, y: int) -> None:
        x = x + x if x >= 0 else -1 - x - x
        y = y + y if y >= 0 else -1 - y - y
        self._graph[x ^ 1].add(y)
        self._graph[y ^ 1].add(x)

    def x_eq_y(self, x: int, y: int) -> None:
        x = x + x if x >= 0 else -1 - x - x
        y = y + y if y >= 0 else -1 - y - y
        self._graph[x].add(y)
        self._graph[x ^ 1].add(y ^ 1)
        self._graph[y].add(x)
        self._graph[y ^ 1].add(x ^ 1)

    def at_most_one(self, x: Iterable[int]) -> None:
        x = [x_i + x_i if x_i >= 0 else -1 - x_i - x_i for x_i in x]
        if len(x) <= 4:
            for x_i, x_j in itertools.permutations(x, 2):
                self._graph[x_i].add(x_j ^ 1)
        else:
            y_i = y0 = len(self._graph)
            for x0, x1 in itertools.pairwise(x):
                self._graph.append((y_i + 2, x1 ^ 1))
                self._graph.append((y_i - 1, x0 ^ 1))
                self._graph[x0].add(y_i)
                self._graph[x1].add(y_i + 1)
                y_i += 2
            self._graph[y0 + 1] = (x[0] ^ 1,)
            self._graph[y_i - 2] = (x[-1] ^ 1,)

    def is_satisfiable(self) -> bool:
        if self._ord is None:
            _, scc_no = graph.strongly_connected_component(self._graph)
            self._ord, self._nord = scc_no[::2], scc_no[1::2]
        return all(x != y for x, y in zip(self._ord, self._nord))

    def find_truth_assignment(self) -> list[bool]:
        if not self.is_satisfiable():
            raise ValueError('No truth assignment exists.')
        return [x < y for x, y, _ in zip(self._ord, self._nord, range(self._n))]

• 2-SAT 참고
• ¬이 붙은 변수들을 add_or_cond 의 인자로 넘기기 위해서는 ~를 붙여서 표현한다. 아래 예시 참고.
• 예시 : x1,x2,x3의 세개의 변수로 이루어진 (x1 ∨ ¬x3) ∧ (¬x2 ∨ x3)을 풀 경우,

  • two_sat = twosat.TwoSat(3)
    two_sat.x_or_y(0, ~2)
    two_sat.x_or_y(~1, 2)
    print(two_sat.is_satisfiable())