add two-sat (#41)

* add two-sat

* fix clippy

* fix doctest

Author: cameroncuster

Cargo.toml

@@ -70,6 +70,10 @@ path = "examples/graphs/lca_yosupo.rs"
 name = "scc"
 path = "examples/graphs/scc.rs"
 
+[[example]]
+name = "two_sat"
+path = "examples/graphs/two_sat.rs"
+
 [[example]]
 name = "binom"
 path = "examples/numbers/binom.rs"

examples/graphs/two_sat.rs

@@ -0,0 +1,40 @@
+// verification-helper: PROBLEM https://judge.yosupo.jp/problem/two_sat
+
+use proconio::input;
+use programming_team_code_rust::graphs::two_sat::TwoSat;
+
+fn main() {
+    input! {
+        _p: String,
+        _cnf: String,
+        n: usize,
+        m: usize,
+        clauses: [(i32, i32, usize); m],
+    }
+
+    let mut ts = TwoSat::new(n);
+
+    for (x, y, _) in clauses {
+        let f = x > 0;
+        let g = y > 0;
+        ts.add_clause(
+            x.unsigned_abs() as usize - 1,
+            f,
+            y.unsigned_abs() as usize - 1,
+            g,
+        );
+    }
+
+    let ans = ts.solve();
+
+    if let Some(ans) = ans {
+        println!("s SATISFIABLE");
+        print!("v");
+        for i in 1..=n {
+            print!(" {}", if ans[i - 1] { i as i32 } else { -(i as i32) });
+        }
+        println!(" 0");
+    } else {
+        println!("s UNSATISFIABLE");
+    }
+}

src/graphs/mod.rs

@@ -5,3 +5,4 @@ mod dfs_order;
 pub mod dijk;
 pub mod lca;
 pub mod scc;
+pub mod two_sat;

src/graphs/two_sat.rs

@@ -0,0 +1,71 @@
+//! # Two Satisfiability
+
+use crate::graphs::scc::get_sccs;
+
+/// # Example
+/// ```
+/// use programming_team_code_rust::graphs::two_sat::TwoSat;
+///
+/// let mut ts = TwoSat::new(2);
+/// ts.add_var();
+/// ts.add_clause(0, false, 1, true);
+/// ts.add_clause(1, false, 2, true);
+/// ts.add_clause(2, false, 0, true);
+/// assert_eq!(ts.solve(), Some(vec![false, false, false]));
+/// ```
+pub struct TwoSat {
+    adj: Vec<Vec<usize>>,
+}
+
+impl TwoSat {
+    /// Create a new instance of TwoSat
+    ///
+    /// # Complexity
+    /// - Time: O(n)
+    /// - Space: O(n)
+    pub fn new(n: usize) -> Self {
+        Self {
+            adj: vec![vec![]; 2 * n],
+        }
+    }
+
+    /// Add a new variable
+    ///
+    /// # Complexity
+    /// - Time: O(1)
+    /// - Space: O(1)
+    pub fn add_var(&mut self) {
+        self.adj.push(vec![]);
+        self.adj.push(vec![]);
+    }
+
+    /// Add a clause to the 2-SAT instance
+    ///
+    /// # Complexity
+    /// - Time: O(1)
+    /// - Space: O(1)
+    pub fn add_clause(&mut self, i: usize, f: bool, j: usize, g: bool) {
+        self.adj[2 * i + !f as usize].push(2 * j + g as usize);
+        self.adj[2 * j + !g as usize].push(2 * i + f as usize);
+    }
+
+    /// Solve the 2-SAT instance
+    /// Returns None if there is no solution, otherwise returns a valid assignment
+    ///
+    /// # Complexity
+    /// - Time: O(n)
+    /// - Space: O(n)
+    pub fn solve(&self) -> Option<Vec<bool>> {
+        let (_, scc_id) = get_sccs(&self.adj);
+        let mut res = vec![false; self.adj.len() / 2];
+        for i in 0..self.adj.len() / 2 {
+            if scc_id[2 * i] == scc_id[2 * i + 1] {
+                return None;
+            }
+            res[i] = scc_id[2 * i] < scc_id[2 * i + 1];
+        }
+        Some(res)
+    }
+
+    // TODO: implement this https://github.com/kth-competitive-programming/kactl/blob/main/content/graph/2sat.h#L43-L54
+}