91 lines
2.7 KiB
Python
91 lines
2.7 KiB
Python
"""2-SAT solver via implication graph + strongly connected components (Tarjan).
|
|
|
|
API:
|
|
solve(n, clauses) -> (sat, assignment)
|
|
- n: number of boolean variables x_0..x_{n-1}.
|
|
- clauses: iterable of (a, b) where each literal is an int in [-n, n]\\{0};
|
|
positive k means x_{k-1} true, negative -k means x_{k-1} false.
|
|
Each clause is (lit1 OR lit2). A unit clause can be (lit, lit) or (lit, 0)? We
|
|
require two literals; pass (lit, lit) for unit.
|
|
- Returns (True, [bool,...]) if satisfiable, else (False, None).
|
|
|
|
A clause (u OR v) becomes implications (!u -> v) and (!v -> u).
|
|
Variable i true -> node 2*i, false -> node 2*i+1 (0-based i).
|
|
Negation of node x is x^1.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
|
|
def _solve(n, clauses):
|
|
N = 2 * n
|
|
adj = [[] for _ in range(N)]
|
|
radj = [[] for _ in range(N)]
|
|
|
|
def node(lit):
|
|
# lit: signed 1-based variable index
|
|
v = abs(lit) - 1
|
|
if lit > 0:
|
|
return 2 * v # true
|
|
return 2 * v + 1 # false
|
|
|
|
for a, b in clauses:
|
|
na, nb = node(a), node(b)
|
|
# (!a -> b) and (!b -> a)
|
|
# negation of node x is x^1
|
|
adj[na ^ 1].append(nb)
|
|
adj[nb ^ 1].append(na)
|
|
radj[nb].append(na ^ 1)
|
|
radj[na].append(nb ^ 1)
|
|
|
|
# Tarjan SCC (iterative to avoid recursion limits)
|
|
index_counter = [0]
|
|
stack = []
|
|
on_stack = [False] * N
|
|
indices = [-1] * N
|
|
lowlink = [0] * N
|
|
comp = [-1] * N
|
|
comp_count = [0]
|
|
|
|
import sys
|
|
sys.setrecursionlimit(max(sys.getrecursionlimit(), N + 100))
|
|
|
|
def strongconnect(v):
|
|
indices[v] = index_counter[0]
|
|
lowlink[v] = index_counter[0]
|
|
index_counter[0] += 1
|
|
stack.append(v)
|
|
on_stack[v] = True
|
|
for w in adj[v]:
|
|
if indices[w] == -1:
|
|
strongconnect(w)
|
|
lowlink[v] = min(lowlink[v], lowlink[w])
|
|
elif on_stack[w]:
|
|
lowlink[v] = min(lowlink[v], indices[w])
|
|
if lowlink[v] == indices[v]:
|
|
while True:
|
|
w = stack.pop()
|
|
on_stack[w] = False
|
|
comp[w] = comp_count[0]
|
|
if w == v:
|
|
break
|
|
comp_count[0] += 1
|
|
|
|
for v in range(N):
|
|
if indices[v] == -1:
|
|
strongconnect(v)
|
|
|
|
for i in range(n):
|
|
if comp[2 * i] == comp[2 * i + 1]:
|
|
return False, None
|
|
|
|
# assignment: for each variable, if comp[true_node] > comp[false_node] then true
|
|
# (in reverse topological order of SCCs; higher comp index = earlier in topo order)
|
|
assign = [False] * n
|
|
for i in range(n):
|
|
assign[i] = comp[2 * i] < comp[2 * i + 1]
|
|
return True, assign
|
|
|
|
|
|
def solve(n, clauses):
|
|
return _solve(n, clauses)
|