diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..570dcb0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +target/ +**/*.rs.bk +Cargo.lock +.idea/ +.vscode/ +*.code-workspace \ No newline at end of file diff --git a/.travis.yml b/.travis.yml new file mode 100644 index 0000000..a99bd14 --- /dev/null +++ b/.travis.yml @@ -0,0 +1,15 @@ +language: rust +rust: + #- 1.58.0 # Version currently supported by Codeforces + - stable + - beta + - nightly +before_script: + - rustup component add clippy +script: + - cargo test --verbose + - cargo clippy -- -D warnings +matrix: + allow_failures: + - rust: nightly + fast_finish: true diff --git a/Cargo.toml b/Cargo.toml index e10786a..ae11d8e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,15 @@ [package] -name = "algorithms" -version = "0.1.0" -authors = ["Aram Ebtekar "] +name = "contest-algorithms" +version = "0.3.1-alpha.1" +authors = ["Aram Ebtekar"] +edition = "2024" -[dependencies] +description = "Common algorithms and data structures for programming contests" +repository = "https://github.com/EbTech/rust-algorithms" +readme = "README.md" +keywords = ["competitive", "programming", "codeforces"] +categories = ["algorithms", "data-structures"] +license = "MIT" + +[badges] +travis-ci = { repository = "EbTech/rust-algorithms", branch = "master" } diff --git a/README.md b/README.md index 84140a0..2add26f 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,15 @@ -# Algorithm Cookbook in Rust +# Contest Algorithms in Rust -A collection of classic data structures and algorithms, emphasizing beauty and clarity over full generality. As such, this should not be viewed as a blackbox *library*, but as a whitebox *cookbook* demonstrating the translation of abstract concepts into executable code. I hope it will be useful to students and educators, as well as competition programmers. +[![Crates.io Version](https://img.shields.io/crates/v/contest-algorithms.svg)](https://crates.io/crates/contest-algorithms) +[![Documentation](https://docs.rs/contest-algorithms/badge.svg)](https://docs.rs/contest-algorithms) +[![license](https://img.shields.io/badge/license-MIT-blue.svg)](https://github.com/bevyengine/bevy/blob/master/LICENSE) +[![Crates.io Downloads](https://img.shields.io/crates/d/contest-algorithms.svg)](https://crates.io/crates/contest-algorithms) +[![Build Status](https://travis-ci.org/EbTech/rust-algorithms.svg?branch=master)](https://travis-ci.org/EbTech/rust-algorithms) +[![Gitter](https://badges.gitter.im/rust-algos/community.svg)](https://gitter.im/rust-algos/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge) -This repository is distributed under the [MIT License](LICENSE). The license text need not be included in contest submissions, though I would appreciate linking back to this repo for others to find. Enjoy! +A collection of classic data structures and algorithms, emphasizing usability, beauty and clarity over full generality. As such, this should be viewed not as a blackbox *library*, but as a whitebox *cookbook* demonstrating the design and implementation of algorithms. I hope it will be useful to students and educators, as well as fans of algorithmic programming contests. + +This repository is distributed under the [MIT License](LICENSE). Contest submissions need not include the license text. Enjoy! ## For Students and Educators @@ -10,22 +17,110 @@ When learning a new algorithm or data structure, it's often helpful to see or pl In addition, the Rust language has outstanding pedagogical attributes. Its compiler acts as a teacher, enforcing strict discipline while pointing to clearer ways to structure one's logic. -## For Competition Programmers +## For Programming Contests + +The original intent of this project was to build a reference for use in programming contests. As a result, it contains algorithms that are frequently useful to have in one's toolkit, with an emphasis on code that is concise and easy to modify under time pressure. + +Most competitive programmers rely on C++ for its fast execution time. However, it's notoriously unsafe, diverting a considerable share of the contestant's time and attention on mistake prevention and debugging. Java is the next most popular choice, offering a little safety at some expense to speed of coding and execution. + +To my delight, I found that Rust eliminates entire classes of bugs, while reducing visual clutter to make the rest easier to spot. And, it's *fast*. There's a learning curve, to be sure. However, a proficient Rust programmer stands to gain a competitive advantage as well as a more pleasant experience! -The original intent of this project was to build a reference for use in programming competitions such as [Codeforces](http://codeforces.com) and the [Google Code Jam](https://code.google.com/codejam). As a result, it contains algorithms that are frequently useful to have in one's toolkit, with an emphasis on making the code concise and easy to modify under time pressure. +Some contest sites and online judges that support Rust: +- [Codeforces](https://codeforces.com) +- [CodeChef](https://www.codechef.com) +- [AtCoder](https://atcoder.jp) +- [Kattis](https://open.kattis.com/help/rust) +- [SPOJ](https://www.spoj.com/) +- [LeetCode](https://leetcode.com/contest) +- [HackerRank](https://www.hackerrank.com/contests) +- [Timus](http://acm.timus.ru/help.aspx?topic=rust) -Most competitive programmers use C/C++ because it allows for fast coding as well as fast execution. However, these languages are notoriously unsafe, wasting a considerable share of the contestant's time and attention on accident prevention and debugging. Java is the next most popular choice, offering a bit of safety at some expense to coding and execution speed. To my delight, I found that Rust provides a lot more safety than Java without the visual clutter, and it's *fast*. A proficient Rust programmer stands to gain a competitive advantage as well as a more pleasant experience! +For help in getting started, you may check out [some of my past submissions](https://codeforces.com/contest/1168/submission/55200038) (requires login). ## Programming Language Advocacy -My other goal is to appeal to developers who feel, as I once did, trapped between the lesser of headaches (e.g., C++ and Java), to raise awareness that *it doesn't have to be this way*. Rather than trying to persuade you with words, this repository aims to show by example and ease the learning curve a bit. See [Jim Blandy's *Why Rust?*](http://www.oreilly.com/programming/free/files/why-rust.pdf) for a brief introduction, or just [dive in](https://www.rust-lang.org)! +My other goal is to appeal to developers who feel limited by ancient (yet still mainstream) programming languages, by demonstrating the power of modern techniques. + +Rather than try to persuade you with words, this repository aims to show by example. If you'd like to learn the language, I recommend [the official book](https://doc.rust-lang.org/book/) or [Programming Rust](https://www.amazon.com/Programming-Rust-Fast-Systems-Development-dp-1492052590/dp/1492052590). + +# Contents + +## [Graphs](src/graph/) + +### [Graph representations](src/graph/mod.rs) + +- Integer index-based adjacency list representation +- Disjoint set union + +### [Elementary graph algorithms](src/graph/util.rs) + +- Euler path and tour +- Kruskal's minimum spanning tree +- Dijkstra's single-source shortest paths +- DFS pre-order traversal + +### [Connected components](src/graph/connectivity.rs) + +- Connected components +- Strongly connected components +- Bridges and 2-edge-connected components +- Articulation points and 2-vertex-connected components +- Topological sort +- 2-SAT solver + +### [Network flows](src/graph/flow.rs) + +- Dinic's blocking maximum flow +- Minimum cut +- Hopcroft-Karp bipartite matching +- Minimum cost maximum flow + +## [Math](src/math/) + +### [Number theory](src/math/mod.rs) + +- Greatest common divisor +- Canonical solution to Bezout's identity +- Miller's primality test + +### [Generic FFT](src/math/fft.rs) + +- Fast Fourier transform +- Number theoretic transform +- Convolution + +### [Arithmetic](src/math/num.rs) + +- Rational numbers +- Complex numbers +- Linear algebra +- Safe modular arithmetic + +## [Ordering and search](src/order.rs) + +- Comparator for `PartialOrd` +- Binary search: drop-in replacements for C++ `lower_bound()`/`upper_bound()` +- Merge and mergesort +- Coordinate compression +- Online convex hull trick (update and query the upper envelope of a set of lines) + +## [Associative range query](src/range_query) + +- Statically allocated binary indexed ARQ tree (a.k.a. generic segtree with lazy propagation) +- Dynamically allocated ARQ tree, optionally sparse and persistent +- Mo's algorithm (a.k.a. query square root decomposition) + +## [Scanner](src/scanner.rs) + +- Utility for reading input data ergonomically +- File and standard I/O examples + +## [String processing](src/string_proc.rs) -## Contents +- Generic trie +- Knuth-Morris-Pratt single-pattern string matching +- Aho-Corasick multi-pattern string matching +- Suffix array: O(n log n) construction using counting sort +- Longest common prefix +- Manacher's linear-time palindrome search -- [Basic graph representations](src/graph/mod.rs): adjacency lists, minimum spanning tree, Euler path, disjoint set union -- [Network flows](src/graph/flow.rs): Dinic's blocking flow, Hopcroft-Karp bipartite matching, min cost max flow -- [Connected components](src/graph/connectivity.rs): 2-edge-, 2-vertex- and strongly connected components, bridges, articulation points, topological sort, 2-SAT -- [Associative range query](src/arq_tree.rs): known colloquially as *segtrees* -- [Math](src/math.rs): Euclid's GCD algorithm, Bezout's identity -- [Scanner](src/scanner.rs): utility for reading input data -- [String processing](src/string_proc.rs): Knuth-Morris-Pratt string matching, Manacher's palindrome search diff --git a/src/arq_tree.rs b/src/arq_tree.rs deleted file mode 100644 index 3fbeb2f..0000000 --- a/src/arq_tree.rs +++ /dev/null @@ -1,122 +0,0 @@ -//! Associative Range Query Tree based on [Al.Cash's compact representation] -//! (http://codeforces.com/blog/entry/18051). - -/// Colloquially known as a "segtree" in the sport programming literature, it -/// represents a sequence of elements a_i (0 <= i < size) from a monoid (M, +) on -/// which we want to support fast range operations: -/// -/// - modify(l, r, f) replaces a_i (l <= i <= r) by f(a_i) for a homomorphism f -/// - query(l, r) returns the aggregate a_l + a_{l+1} + ... + a_r -/// -/// To customize, simply change the commented lines. -/// In this example, we chose to support range sum queries and range constant -/// assignments. Since constant assignment f_c(a) = c is not a homomorphism over -/// integers, we have to augment the monoid type, using the 2D vector (a_i, 1) -/// instead of a_i. You may check that f_c((a, s)) = (c*s, s) is a homomorphism. -pub struct ArqTree { - d: Vec>, - t: Vec, - s: Vec, -} - -impl ArqTree { - /// Initializes a sequence of identity elements. - pub fn new(size: usize) -> Self { - let mut s = vec![1; 2 * size]; - for i in (0..size).rev() { - s[i] = s[i << 1] + s[i << 1 | 1]; - } - Self { - d: vec![None; size], - t: vec![0; 2 * size], // monoid identity - s: s, - } - } - - fn apply(&mut self, p: usize, f: i64) { - self.t[p] = f * self.s[p]; // hom application - if p < self.d.len() { - self.d[p] = Some(f); // hom composition - } - } - - fn push(&mut self, p: usize) { - for s in (1..32).rev() { - let i = p >> s; - if let Some(f) = self.d[i] { - self.apply(i << 1, f); - self.apply(i << 1 | 1, f); - self.d[i] = None; - } - } - } - - fn pull(&mut self, mut p: usize) { - while p > 1 { - p >>= 1; - if self.d[p] == None { - self.t[p] = self.t[p << 1] + self.t[p << 1 | 1]; // monoid op - } - } - } - - /// Performs the homomorphism f on all entries from l to r, inclusive. - pub fn modify(&mut self, mut l: usize, mut r: usize, f: i64) { - l += self.d.len(); - r += self.d.len(); - let (l0, r0) = (l, r); - self.push(l0); - self.push(r0); - while l <= r { - if l & 1 == 1 { - self.apply(l, f); - l += 1; - } - if r & 1 == 0 { - self.apply(r, f); - r -= 1; - } - l >>= 1; - r >>= 1; - } - self.pull(l0); - self.pull(r0); - } - - /// Returns the aggregate range query on all entries from l to r, inclusive. - pub fn query(&mut self, mut l: usize, mut r: usize) -> i64 { - l += self.d.len(); - r += self.d.len(); - self.push(l); - self.push(r); - let mut res = 0; // monoid identity - while l <= r { - if l & 1 == 1 { - res = res + self.t[l]; // monoid op - l += 1; - } - if r & 1 == 0 { - res = self.t[r] + res; // monoid op - r -= 1; - } - l >>= 1; - r >>= 1; - } - res - } -} - -#[cfg(test)] -mod test { - use super::*; - - #[test] - fn test_arq_tree() { - let mut arq = ArqTree::new(10); - - arq.modify(1, 3, 10); - arq.modify(3, 5, 1); - - assert_eq!(arq.query(0, 9), 23); - } -} diff --git a/src/caching.rs b/src/caching.rs new file mode 100644 index 0000000..ed22656 --- /dev/null +++ b/src/caching.rs @@ -0,0 +1,134 @@ +//! Basic Cacher struct which stores a closure and a hashmap. +//! The hashmap stores key value pairs representing previous +//! function calls. +//! +//! When the Cacher function is run, it first does a lookup +//! to see if the value has already been calculated. If it has, +//! it returns that value. If it hasn't, it calculates the value, +//! adds it to the hashmap, and returns it. + +use std::collections::HashMap; + +/// The Cacher struct (Memoization) stores a function and a Hashmap. +/// The HashMap keeps track of previous input and output for the function so +/// that it only ever has to be called once per input. Use for expensive functions. +pub struct Cacher +where + F: Fn(U) -> V, + U: std::cmp::Eq + std::hash::Hash + Copy, + V: Copy, +{ + calculation: F, + values: HashMap, +} + +impl Cacher +where + F: Fn(U) -> V, + U: std::cmp::Eq + std::hash::Hash + Copy, + V: Copy, +{ + /// Constuctor for the Casher + /// # Examples + /// ``` + /// # use contest_algorithms::caching::Cacher; + /// let mut squared = Cacher::new(|n: u32| n*n); + /// ``` + pub fn new(calculation: F) -> Cacher { + Cacher { + calculation, + values: HashMap::new(), + } + } + + /// Performs a lookup into the HashMap to see if the value has already + /// been calculated. If it has, returns the value. If it has not, + /// calls the function, stores the value, then returns the value. + /// # Examples + /// ``` + /// # use contest_algorithms::caching::Cacher; + /// let mut squared = Cacher::new(|n: u32| n*n); + /// + /// // This is where we call the function + /// let sixteen = squared.call(4); + /// ``` + // TODO: whenever Rust's Entry API gains the ability to take ownership of + // arg only when necessary, this method should follow the same practice. + // Also, Cacher should implement Fn(U)->V once this is possible. + pub fn call(&mut self, arg: U) -> V { + let calc = &self.calculation; + *self.values.entry(arg).or_insert_with_key(|&key| calc(key)) + } + + /// Calls the function without performing a lookup and replaces + /// the old return value with the new one, and returns it. + /// Potentially useful if the function reads from a file or RNG + /// whose state may have changed. + // TODO: if there's state, FnMut seems more appropriate. + pub fn call_and_replace(&mut self, arg: U) -> V { + let new_val = (self.calculation)(arg); + self.values.insert(arg, new_val); + new_val + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_cacher_basically_works() { + let mut word_len = Cacher::new(|word: &str| word.len()); + let hello = word_len.call("hello"); + + // Test function returns correctly + assert_eq!(hello, 5); + + // Test HashMap is correct length + assert_eq!(word_len.values.len(), 1); + + // Test HashMap has correct value after one insert + let mut test_map = HashMap::new(); + test_map.insert("hello", 5); + assert_eq!(word_len.values, test_map); + + // Test HashMap has correct value after duplicate insert + word_len.call("hello"); + assert_eq!(word_len.values, test_map); + + // Test HashMap has correct values after unique input + word_len.call("wazzup"); + test_map.insert("wazzup", 6); + assert_eq!(word_len.values, test_map); + } + + #[test] + fn test_cacher_speed() { + // Simulate a function that takes 1 second to complete + let mut func = Cacher::new(|x| { + std::thread::sleep(std::time::Duration::from_millis(100)); + x * x + }); + + // Would take 10 minutes without caching + for _ in 0..6000 { + assert_eq!(25, func.call(5)); + } + } + + #[test] + fn test_call_and_replace() { + use std::time::Instant; + + let mut func = Cacher::new(|_param: usize| Instant::now()); + let first_instant = func.call(0); + let lookup_instant = func.call(0); + + assert_eq!(first_instant, lookup_instant); + assert_eq!(1, func.values.len()); + + let second_instant = func.call_and_replace(0); + assert_eq!(1, func.values.len()); + assert_ne!(second_instant, lookup_instant); + } +} diff --git a/src/graph/connectivity.rs b/src/graph/connectivity.rs index d1246ca..aaf854c 100644 --- a/src/graph/connectivity.rs +++ b/src/graph/connectivity.rs @@ -1,6 +1,40 @@ //! Graph connectivity structures. -use graph::Graph; -use std::cmp::min; +use super::Graph; + +/// Helper struct that carries data needed for the depth-first searches in +/// ConnectivityGraph's constructor. +struct ConnectivityData { + time: usize, + vis: Box<[usize]>, + low: Box<[usize]>, + v_stack: Vec, + e_stack: Vec, +} + +impl ConnectivityData { + fn new(num_v: usize) -> Self { + Self { + time: 0, + vis: vec![0; num_v].into_boxed_slice(), + low: vec![0; num_v].into_boxed_slice(), + v_stack: vec![], + e_stack: vec![], + } + } + + fn visit(&mut self, u: usize) { + self.time += 1; + self.vis[u] = self.time; + self.low[u] = self.time; + self.v_stack.push(u); + } + + fn lower(&mut self, u: usize, val: usize) { + if self.low[u] > val { + self.low[u] = val + } + } +} /// Represents the decomposition of a graph into any of its constituent parts: /// @@ -9,75 +43,62 @@ use std::cmp::min; /// - 2-edge-connected components (2ECC), /// - 2-vertex-connected components (2VCC) /// -/// Multiple-edges and self-loops should be correctly handled. +/// Multiple-edges and self-loops are correctly handled. pub struct ConnectivityGraph<'a> { - pub graph: &'a Graph, // Immutable graph, frozen for lifetime of this obj. - pub cc: Vec, // Stores id of a vertex's CC, SCC or 2ECC. - pub vcc: Vec, // Stores id of an edge's 2VCC. + /// Immutable graph, frozen for the lifetime of the ConnectivityGraph object. + pub graph: &'a Graph, + /// ID of a vertex's CC, SCC or 2ECC, whichever applies. Range 1 to num_cc. + pub cc: Vec, + /// ID of an edge's 2VCC, where applicable. Ranges from 1 to num_vcc. + pub vcc: Vec, + /// Total number of CCs, SCCs or 2ECCs, whichever applies. pub num_cc: usize, + /// Total number of 2VCCs, where applicable. pub num_vcc: usize, } impl<'a> ConnectivityGraph<'a> { - /// Computes the SCCs of a directed graph in reverse topological order, or - /// the 2ECCs/2VCCS of an undirected graph. Can also get CCs by passing an - /// undirected graph using `is_directed == true`. + /// Computes CCs (connected components), SCCs (strongly connected + /// components), 2ECCs (2-edge-connected components), and/or 2VCCs + /// (2-vertex-connected components), depending on the parameter and graph: + /// - is_directed == true on directed graph: SCCs in rev-topological order + /// - is_directed == true on undirected graph: CCs + /// - is_directed == false on undirected graph: 2ECCs and 2VCCs + /// - is_directed == false on directed graph: undefined behavior pub fn new(graph: &'a Graph, is_directed: bool) -> Self { let mut connect = Self { - graph: graph, + graph, cc: vec![0; graph.num_v()], vcc: vec![0; graph.num_e()], num_cc: 0, num_vcc: 0, }; - let mut t = 0; - let mut vis = vec![0; graph.num_v()]; - let mut low = vec![0; graph.num_v()]; - let mut v_stack = Vec::new(); - let mut e_stack = Vec::new(); + let mut data = ConnectivityData::new(graph.num_v()); for u in 0..graph.num_v() { - if vis[u] == 0 { + if data.vis[u] == 0 { if is_directed { - connect.scc(u, &mut t, &mut vis, &mut low, &mut v_stack); + connect.scc(&mut data, u); } else { - connect.bcc( - u, - graph.num_e() + 1, - &mut t, - &mut vis, - &mut low, - &mut v_stack, - &mut e_stack, - ); + connect.bcc(&mut data, u, graph.num_e() + 1); } } } connect } - fn scc( - &mut self, - u: usize, - t: &mut usize, - vis: &mut [usize], - low: &mut [usize], - v_stack: &mut Vec, - ) { - *t += 1; - vis[u] = *t; - low[u] = *t; - v_stack.push(u); + fn scc(&mut self, data: &mut ConnectivityData, u: usize) { + data.visit(u); for (_, v) in self.graph.adj_list(u) { - if vis[v] == 0 { - self.scc(v, t, vis, low, v_stack); + if data.vis[v] == 0 { + self.scc(data, v); } if self.cc[v] == 0 { - low[u] = min(low[u], low[v]); + data.lower(u, data.low[v]); } } - if vis[u] == low[u] { + if data.vis[u] == data.low[u] { self.num_cc += 1; - while let Some(v) = v_stack.pop() { + while let Some(v) = data.v_stack.pop() { self.cc[v] = self.num_cc; if v == u { break; @@ -102,36 +123,25 @@ impl<'a> ConnectivityGraph<'a> { .collect() } - /// Gets the vertices of a directed acyclic graph (DAG) in topological order. + /// Gets the vertices of a graph according to a topological order of the + /// strongly connected components. Most often used on DAGs. pub fn topological_sort(&self) -> Vec { let mut vertices = (0..self.graph.num_v()).collect::>(); - vertices.sort_by_key(|&u| self.num_cc - self.cc[u]); + vertices.sort_unstable_by_key(|&u| self.num_cc - self.cc[u]); vertices } - fn bcc( - &mut self, - u: usize, - par: usize, - t: &mut usize, - vis: &mut [usize], - low: &mut [usize], - v_stack: &mut Vec, - e_stack: &mut Vec, - ) { - *t += 1; - vis[u] = *t; - low[u] = *t; - v_stack.push(u); + fn bcc(&mut self, data: &mut ConnectivityData, u: usize, par: usize) { + data.visit(u); for (e, v) in self.graph.adj_list(u) { - if vis[v] == 0 { - e_stack.push(e); - self.bcc(v, e, t, vis, low, v_stack, e_stack); - low[u] = min(low[u], low[v]); - if vis[u] <= low[v] { + if data.vis[v] == 0 { + data.e_stack.push(e); + self.bcc(data, v, e); + data.lower(u, data.low[v]); + if data.vis[u] <= data.low[v] { // u is a cut vertex unless it's a one-child root self.num_vcc += 1; - while let Some(top_e) = e_stack.pop() { + while let Some(top_e) = data.e_stack.pop() { self.vcc[top_e] = self.num_vcc; self.vcc[top_e ^ 1] = self.num_vcc; if e ^ top_e <= 1 { @@ -139,9 +149,9 @@ impl<'a> ConnectivityGraph<'a> { } } } - } else if vis[v] < vis[u] && e ^ par != 1 { - low[u] = min(low[u], vis[v]); - e_stack.push(e); + } else if data.vis[v] < data.vis[u] && e ^ par != 1 { + data.lower(u, data.vis[v]); + data.e_stack.push(e); } else if v == u { // e is a self-loop self.num_vcc += 1; @@ -149,10 +159,10 @@ impl<'a> ConnectivityGraph<'a> { self.vcc[e ^ 1] = self.num_vcc; } } - if vis[u] == low[u] { + if data.vis[u] == data.low[u] { // par is a cut edge unless par==-1 self.num_cc += 1; - while let Some(v) = v_stack.pop() { + while let Some(v) = data.v_stack.pop() { self.cc[v] = self.num_cc; if v == u { break; @@ -164,9 +174,9 @@ impl<'a> ConnectivityGraph<'a> { /// In an undirected graph, determines whether u is an articulation vertex. pub fn is_cut_vertex(&self, u: usize) -> bool { if let Some(first_e) = self.graph.first[u] { - self.graph.adj_list(u).any(|(e, _)| { - self.vcc[first_e] != self.vcc[e] - }) + self.graph + .adj_list(u) + .any(|(e, _)| self.vcc[first_e] != self.vcc[e]) } else { false } @@ -186,7 +196,8 @@ mod test { #[test] fn test_toposort() { - let mut graph = Graph::new(4, 4); + let mut graph = Graph::new(4, 5); + graph.add_edge(0, 0); graph.add_edge(0, 2); graph.add_edge(3, 2); graph.add_edge(3, 1); diff --git a/src/graph/flow.rs b/src/graph/flow.rs index 0f17bf6..41af65a 100644 --- a/src/graph/flow.rs +++ b/src/graph/flow.rs @@ -1,64 +1,73 @@ -//! Maximum flows and minimum cuts. -use graph::{Graph, AdjListIterator}; -use std::cmp::min; -const INF: i64 = 0x3f3f3f3f; +//! Maximum flows, matchings, and minimum cuts. +use super::{AdjListIterator, Graph}; /// Representation of a network flow problem with (optional) costs. pub struct FlowGraph { - pub graph: Graph, // Owned graph, controlled by this FlowGraph object. + /// Owned graph, managed by this FlowGraph object. + pub graph: Graph, + /// Edge capacities. pub cap: Vec, + /// Edge cost per unit flow. pub cost: Vec, } impl FlowGraph { + /// An upper limit to the flow. + const INF: i64 = i64::MAX; + /// Initializes an flow network with vmax vertices and no edges. - pub fn new(vmax: usize, emax: usize) -> Self { + pub fn new(vmax: usize, emax_hint: usize) -> Self { Self { - graph: Graph::new(vmax, 2 * emax), - cap: Vec::with_capacity(2 * emax), - cost: Vec::with_capacity(2 * emax), + graph: Graph::new(vmax, 2 * emax_hint), + cap: Vec::with_capacity(2 * emax_hint), + cost: Vec::with_capacity(2 * emax_hint), } } - /// Adds an edge with specified capacity and cost. The reverse edge is also - /// added for residual graph computation, but has zero capacity. - pub fn add_edge(&mut self, u: usize, v: usize, cap: i64, cost: i64) { + /// Adds an edge with specified directional capacities and cost per unit of + /// flow. If only forward flow is allowed, rcap should be zero. + pub fn add_edge(&mut self, u: usize, v: usize, cap: i64, rcap: i64, cost: i64) { self.cap.push(cap); - self.cap.push(0); + self.cap.push(rcap); self.cost.push(cost); self.cost.push(-cost); self.graph.add_undirected_edge(u, v); } - /// Dinic's maximum flow / Hopcroft-Karp maximum bipartite matching: + /// Dinic's algorithm to find the maximum flow from s to t where s != t. + /// Generalizes the Hopcroft-Karp maximum bipartite matching algorithm. /// V^2E in general, min(V^(2/3),sqrt(E))E when all edges are unit capacity, /// sqrt(V)E when all vertices are unit capacity as in bipartite graphs. - pub fn dinic(&self, s: usize, t: usize) -> i64 { + /// + /// # Panics + /// + /// Panics if the maximum flow is 2^63 or larger. + pub fn dinic(&self, s: usize, t: usize) -> (i64, Vec) { let mut flow = vec![0; self.graph.num_e()]; let mut max_flow = 0; loop { let dist = self.dinic_search(s, &flow); - if dist[t] == INF { + if dist[t] == Self::INF { break; } // Keep track of adjacency lists to avoid revisiting blocked edges. let mut adj_iters = (0..self.graph.num_v()) .map(|u| self.graph.adj_list(u).peekable()) .collect::>(); - max_flow += self.dinic_augment(s, t, INF, &dist, &mut adj_iters, &mut flow); + max_flow += self.dinic_augment(s, t, Self::INF, &dist, &mut adj_iters, &mut flow); } - max_flow + (max_flow, flow) } // Compute BFS distances to restrict attention to shortest path edges. fn dinic_search(&self, s: usize, flow: &[i64]) -> Vec { - let mut dist = vec![INF; self.graph.num_v()]; + let mut dist = vec![Self::INF; self.graph.num_v()]; let mut q = ::std::collections::VecDeque::new(); dist[s] = 0; q.push_back(s); while let Some(u) = q.pop_front() { for (e, v) in self.graph.adj_list(u) { - if dist[v] == INF && flow[e] < self.cap[e] { + if dist[v] == Self::INF && flow[e] < self.cap[e] { dist[v] = dist[u] + 1; q.push_back(v); } @@ -83,7 +92,7 @@ impl FlowGraph { let mut df = 0; while let Some(&(e, v)) = adj[u].peek() { - let rem_cap = min(self.cap[e] - flow[e], f - df); + let rem_cap = (self.cap[e] - flow[e]).min(f - df); if rem_cap > 0 && dist[v] == dist[u] + 1 { let cf = self.dinic_augment(v, t, rem_cap, dist, adj, flow); flow[e] += cf; @@ -96,7 +105,7 @@ impl FlowGraph { // The current edge is either saturated or blocked. adj[u].next(); } - return df; + df } /// After running maximum flow, use this to recover the dual minimum cut. @@ -105,13 +114,18 @@ impl FlowGraph { .filter(|&e| { let u = self.graph.endp[e ^ 1]; let v = self.graph.endp[e]; - dist[u] < INF && dist[v] == INF + dist[u] < Self::INF && dist[v] == Self::INF }) .collect() } - /// Minimum cost maximum flow, assuming no negative-cost cycles. - pub fn mcf(&self, s: usize, t: usize) -> (i64, i64) { + /// Among all s-t maximum flows, finds one with minimum cost, assuming + /// s != t and no negative-cost cycles. + /// + /// # Panics + /// + /// Panics if the flow or cost overflow a 64-bit signed integer. + pub fn mcf(&self, s: usize, t: usize) -> (i64, i64, Vec) { let mut pot = vec![0; self.graph.num_v()]; // Bellman-Ford deals with negative-cost edges at initialization. @@ -120,7 +134,7 @@ impl FlowGraph { if self.cap[e] > 0 { let u = self.graph.endp[e ^ 1]; let v = self.graph.endp[e]; - pot[v] = min(pot[v], pot[u] + self.cost[e]); + pot[v] = pot[v].min(pot[u] + self.cost[e]); } } } @@ -129,26 +143,26 @@ impl FlowGraph { let (mut min_cost, mut max_flow) = (0, 0); loop { let par = self.mcf_search(s, &flow, &mut pot); - if par[t] == None { + if par[t].is_none() { break; } let (dc, df) = self.mcf_augment(t, &par, &mut flow); min_cost += dc; max_flow += df; } - (min_cost, max_flow) + (min_cost, max_flow, flow) } - // Maintains Johnson's potentials to prevent negative-weight residual edges. - // This enables running Dijkstra instead of the slower Bellman-Ford. + // Maintains Johnson's potentials to prevent negative-cost residual edges. + // This allows running Dijkstra instead of the slower Bellman-Ford. fn mcf_search(&self, s: usize, flow: &[i64], pot: &mut [i64]) -> Vec> { let mut vis = vec![false; self.graph.num_v()]; - let mut dist = vec![INF; self.graph.num_v()]; + let mut dist = vec![Self::INF; self.graph.num_v()]; let mut par = vec![None; self.graph.num_v()]; dist[s] = 0; while let Some(u) = (0..self.graph.num_v()) - .filter(|&u| !vis[u] && dist[u] < INF) + .filter(|&u| !vis[u] && dist[u] < Self::INF) .min_by_key(|&u| dist[u] - pot[u]) { vis[u] = true; @@ -165,10 +179,10 @@ impl FlowGraph { // Pushes flow along an augmenting path of minimum cost. fn mcf_augment(&self, t: usize, par: &[Option], flow: &mut [i64]) -> (i64, i64) { - let (mut dc, mut df) = (0, INF); + let (mut dc, mut df) = (0, Self::INF); let mut u = t; while let Some(e) = par[u] { - df = min(df, self.cap[e] - flow[e]); + df = df.min(self.cap[e] - flow[e]); u = self.graph.endp[e ^ 1]; } u = t; @@ -189,23 +203,73 @@ mod test { #[test] fn test_basic_flow() { let mut graph = FlowGraph::new(3, 2); - graph.add_edge(0, 1, 4, 1); - graph.add_edge(1, 2, 3, 1); + graph.add_edge(0, 1, 4, 0, 0); + graph.add_edge(1, 2, 3, 0, 0); - let flow = graph.dinic(0, 2); + let flow = graph.dinic(0, 2).0; assert_eq!(flow, 3); } #[test] fn test_min_cost_flow() { let mut graph = FlowGraph::new(4, 4); - graph.add_edge(0, 1, 10, -10); - graph.add_edge(1, 2, 7, 8); - graph.add_edge(2, 3, 7, 8); - graph.add_edge(1, 3, 7, 10); + graph.add_edge(0, 1, 10, 0, -10); + graph.add_edge(1, 2, 7, 0, 8); + graph.add_edge(2, 3, 7, 0, 8); + graph.add_edge(1, 3, 7, 0, 10); - let (cost, flow) = graph.mcf(0, 3); + let (cost, flow, _) = graph.mcf(0, 3); assert_eq!(cost, 18); assert_eq!(flow, 10); } + + #[test] + fn test_max_matching() { + let mut graph = FlowGraph::new(14, 4); + + let source = 0; + let sink = 13; + + //Vertex indices of "left hand side" of bipartite graph go from [left_start, right_start) + let left_start = 1; + //Vertex indices of "right hand side" of bipartite graph go from [right_start, sink) + let right_start = 7; + + //Initialize source / sink connections; both left & right have 6 nodes + for lhs_vertex in left_start..left_start + 6 { + graph.add_edge(source, lhs_vertex, 1, 0, 0); + } + + for rhs_vertex in right_start..right_start + 6 { + graph.add_edge(rhs_vertex, sink, 1, 0, 0); + } + + graph.add_edge(left_start + 0, right_start + 1, 1, 0, 0); + graph.add_edge(left_start + 0, right_start + 2, 1, 0, 0); + graph.add_edge(left_start + 2, right_start + 0, 1, 0, 0); + graph.add_edge(left_start + 2, right_start + 3, 1, 0, 0); + graph.add_edge(left_start + 3, right_start + 2, 1, 0, 0); + graph.add_edge(left_start + 4, right_start + 2, 1, 0, 0); + graph.add_edge(left_start + 4, right_start + 3, 1, 0, 0); + graph.add_edge(left_start + 5, right_start + 5, 1, 0, 0); + + let (flow_amt, flow) = graph.dinic(source, sink); + assert_eq!(flow_amt, 5); + + //L->R edges in maximum matching + let left_right_edges = flow + .into_iter() + .enumerate() + .filter(|&(_e, f)| f > 0) + //map to u->v + .map(|(e, _f)| (graph.graph.endp[e ^ 1], graph.graph.endp[e])) + //leave out source and sink nodes + .filter(|&(u, v)| u != source && v != sink) + .collect::>(); + + assert_eq!( + left_right_edges, + vec![(1, 8), (3, 7), (4, 9), (5, 10), (6, 12)] + ); + } } diff --git a/src/graph/mod.rs b/src/graph/mod.rs index 6ed10ef..f5a4c9e 100644 --- a/src/graph/mod.rs +++ b/src/graph/mod.rs @@ -1,6 +1,11 @@ -//! Basic graph library without explicit support for deletion. -pub mod flow; +//! Basic graph module without explicit support for deletion. +//! +//! # Panics +//! +//! All methods will panic if given an out-of-bounds element index. pub mod connectivity; +pub mod flow; +pub mod util; /// Represents a union of disjoint sets. Each set's elements are arranged in a /// tree, whose root is the set's representative. @@ -11,7 +16,9 @@ pub struct DisjointSets { impl DisjointSets { /// Initializes disjoint sets containing one element each. pub fn new(size: usize) -> Self { - Self { parent: (0..size).collect() } + Self { + parent: (0..size).collect(), + } } /// Finds the set's representative. Do path compression along the way to make @@ -33,21 +40,26 @@ impl DisjointSets { } } -/// A compact graph representation. +/// A compact graph representation. Edges are numbered in order of insertion. +/// Each adjacency list consists of all edges pointing out from a given vertex. pub struct Graph { - pub first: Vec>, - pub next: Vec>, - pub endp: Vec, + /// Maps a vertex id to the first edge in its adjacency list. + first: Vec>, + /// Maps an edge id to the next edge in the same adjacency list. + next: Vec>, + /// Maps an edge id to the vertex that it points to. + endp: Vec, } impl Graph { - /// Initializes a graph with vmax vertices and no edges. For best efficiency, - /// emax should be a tight upper bound on the number of edges to insert. - pub fn new(vmax: usize, emax: usize) -> Self { + /// Initializes a graph with vmax vertices and no edges. To reduce + /// unnecessary allocations, emax_hint should be close to the number of + /// edges that will be inserted. + pub fn new(vmax: usize, emax_hint: usize) -> Self { Self { first: vec![None; vmax], - next: Vec::with_capacity(emax), - endp: Vec::with_capacity(emax), + next: Vec::with_capacity(emax_hint), + endp: Vec::with_capacity(emax_hint), } } @@ -75,9 +87,10 @@ impl Graph { self.add_edge(v, u); } - /// If we think of each even-numbered vertex as a variable, and its successor - /// as its negation, then we can build the implication graph corresponding - /// to any 2-CNF formula. Note that u||v == !u -> v == !v -> u. + /// If we think of each even-numbered vertex as a variable, and its + /// odd-numbered successor as its negation, then we can build the + /// implication graph corresponding to any 2-CNF formula. + /// Note that u||v == !u -> v == !v -> u. pub fn add_two_sat_clause(&mut self, u: usize, v: usize) { self.add_edge(u ^ 1, v); self.add_edge(v ^ 1, u); @@ -90,43 +103,6 @@ impl Graph { next_e: self.first[u], } } - - /// Finds the sequence of edges in an Euler path starting from u, assuming it - /// exists and that the graph is directed. To extend this to undirected - /// graphs, keep track of a visited array to skip the reverse edge. - pub fn euler_path(&self, u: usize) -> Vec { - let mut adj_iters = (0..self.num_v()) - .map(|u| self.adj_list(u)) - .collect::>(); - let mut edges = Vec::with_capacity(self.num_e()); - self.euler_recurse(u, &mut adj_iters, &mut edges); - edges.reverse(); - edges - } - - // Helper function used by euler_path. Note that we can't consume the - // adjacency list in a for loop because recursive calls may need it. - fn euler_recurse(&self, u: usize, adj: &mut [AdjListIterator], edges: &mut Vec) { - while let Some((e, v)) = adj[u].next() { - self.euler_recurse(v, adj, edges); - edges.push(e); - } - } - - /// Kruskal's minimum spanning tree algorithm on an undirected graph. - pub fn min_spanning_tree(&self, weights: &[i64]) -> Vec { - assert_eq!(self.num_e(), 2 * weights.len()); - let mut edges = (0..weights.len()).collect::>(); - edges.sort_by_key(|&e| weights[e]); - - let mut components = DisjointSets::new(self.num_v()); - edges - .into_iter() - .filter(|&e| { - components.merge(self.endp[2 * e], self.endp[2 * e + 1]) - }) - .collect() - } } /// An iterator for convenient adjacency list traversal. @@ -135,7 +111,7 @@ pub struct AdjListIterator<'a> { next_e: Option, } -impl<'a> Iterator for AdjListIterator<'a> { +impl Iterator for AdjListIterator<'_> { type Item = (usize, usize); /// Produces an outgoing edge and vertex. @@ -153,27 +129,19 @@ mod test { use super::*; #[test] - fn test_euler() { - let mut graph = Graph::new(3, 4); - graph.add_edge(0, 1); - graph.add_edge(1, 0); + fn test_adj_list() { + let mut graph = Graph::new(5, 6); + graph.add_edge(2, 3); + graph.add_edge(2, 4); + graph.add_edge(4, 1); graph.add_edge(1, 2); - graph.add_edge(2, 1); + graph.add_undirected_edge(0, 2); - assert_eq!(graph.euler_path(0), vec![0, 2, 3, 1]); - } + let adj = graph.adj_list(2).collect::>(); - #[test] - fn test_min_spanning_tree() { - let mut graph = Graph::new(3, 3); - graph.add_undirected_edge(0, 1); - graph.add_undirected_edge(1, 2); - graph.add_undirected_edge(2, 0); - let weights = [7, 3, 5]; - - let mst = graph.min_spanning_tree(&weights); - let mst_cost = mst.iter().map(|&e| weights[e]).sum::(); - assert_eq!(mst, vec![1, 2]); - assert_eq!(mst_cost, 8); + assert_eq!(adj, vec![(5, 0), (1, 4), (0, 3)]); + for (e, v) in adj { + assert_eq!(v, graph.endp[e]); + } } } diff --git a/src/graph/util.rs b/src/graph/util.rs new file mode 100644 index 0000000..bf38f19 --- /dev/null +++ b/src/graph/util.rs @@ -0,0 +1,207 @@ +use super::{DisjointSets, Graph}; +use crate::graph::AdjListIterator; +use std::cmp::Reverse; + +impl Graph { + /// Finds the sequence of edges in an Euler path starting from u, assuming + /// it exists and that the graph is directed. Undefined behavior if this + /// precondition is violated. To extend this to undirected graphs, maintain + /// a visited array to skip the reverse edge. + pub fn euler_path(&self, u: usize) -> Vec { + let mut adj_iters = (0..self.num_v()) + .map(|u| self.adj_list(u)) + .collect::>(); + let mut edges = Vec::with_capacity(self.num_e()); + Self::euler_recurse(u, &mut adj_iters, &mut edges); + edges.reverse(); + edges + } + + // Helper function used by euler_path. Note that we can't use a for-loop + // that would consume the adjacency list as recursive calls may need it. + fn euler_recurse(u: usize, adj: &mut [AdjListIterator], edges: &mut Vec) { + while let Some((e, v)) = adj[u].next() { + Self::euler_recurse(v, adj, edges); + edges.push(e); + } + } + + /// Kruskal's minimum spanning tree algorithm on an undirected graph. + pub fn min_spanning_tree(&self, weights: &[i64]) -> Vec { + assert_eq!(self.num_e(), 2 * weights.len()); + let mut edges = (0..weights.len()).collect::>(); + edges.sort_unstable_by_key(|&e| weights[e]); + + let mut components = DisjointSets::new(self.num_v()); + edges + .into_iter() + .filter(|&e| components.merge(self.endp[2 * e], self.endp[2 * e + 1])) + .collect() + } + + // Single-source shortest paths on a directed graph with non-negative weights + pub fn dijkstra(&self, weights: &[u64], u: usize) -> Vec { + assert_eq!(self.num_e(), weights.len()); + let mut dist = vec![u64::MAX; weights.len()]; + let mut heap = std::collections::BinaryHeap::new(); + + dist[u] = 0; + heap.push((Reverse(0), 0)); + while let Some((Reverse(dist_u), u)) = heap.pop() { + if dist[u] == dist_u { + for (e, v) in self.adj_list(u) { + let dist_v = dist_u + weights[e]; + if dist[v] > dist_v { + dist[v] = dist_v; + heap.push((Reverse(dist_v), v)); + } + } + } + } + dist + } + + pub fn dfs(&self, root: usize) -> DfsIterator { + let mut visited = vec![false; self.num_v()]; + visited[root] = true; + let adj_iters = (0..self.num_v()) + .map(|u| self.adj_list(u)) + .collect::>(); + + DfsIterator { + visited, + stack: vec![root], + adj_iters, + } + } +} + +pub struct DfsIterator<'a> { + visited: Vec, + stack: Vec, + adj_iters: Vec>, +} + +impl Iterator for DfsIterator<'_> { + type Item = (usize, usize); + + /// Returns next edge and vertex in the depth-first traversal + // Refs: https://www.geeksforgeeks.org/iterative-depth-first-traversal/ + // https://en.wikipedia.org/wiki/Depth-first_search + fn next(&mut self) -> Option { + loop { + let &u = self.stack.last()?; + for (e, v) in self.adj_iters[u].by_ref() { + if !self.visited[v] { + self.visited[v] = true; + self.stack.push(v); + return Some((e, v)); + } + } + self.stack.pop(); + } + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_euler() { + let mut graph = Graph::new(3, 4); + graph.add_edge(0, 1); + graph.add_edge(1, 0); + graph.add_edge(1, 2); + graph.add_edge(2, 1); + + assert_eq!(graph.euler_path(0), vec![0, 2, 3, 1]); + } + + #[test] + fn test_min_spanning_tree() { + let mut graph = Graph::new(3, 6); + graph.add_undirected_edge(0, 1); + graph.add_undirected_edge(1, 2); + graph.add_undirected_edge(2, 0); + let weights = [7, 3, 5]; + + let mst = graph.min_spanning_tree(&weights); + let mst_cost = mst.iter().map(|&e| weights[e]).sum::(); + assert_eq!(mst, vec![1, 2]); + assert_eq!(mst_cost, 8); + } + + #[test] + fn test_dijkstra() { + let mut graph = Graph::new(3, 3); + graph.add_edge(0, 1); + graph.add_edge(1, 2); + graph.add_edge(2, 0); + let weights = [7, 3, 5]; + + let dist = graph.dijkstra(&weights, 0); + assert_eq!(dist, vec![0, 7, 10]); + } + + #[test] + fn test_dfs() { + let mut graph = Graph::new(4, 6); + graph.add_edge(0, 2); + graph.add_edge(2, 0); + graph.add_edge(1, 2); + graph.add_edge(0, 1); + graph.add_edge(3, 3); + graph.add_edge(2, 3); + + let dfs_root = 2; + let dfs_traversal = std::iter::once(dfs_root) + .chain(graph.dfs(dfs_root).map(|(_, v)| v)) + .collect::>(); + + assert_eq!(dfs_traversal, vec![2, 3, 0, 1]); + } + + #[test] + fn test_dfs2() { + let mut graph = Graph::new(5, 6); + graph.add_edge(0, 2); + graph.add_edge(2, 1); + graph.add_edge(1, 0); + graph.add_edge(0, 3); + graph.add_edge(3, 4); + graph.add_edge(4, 0); + + let dfs_root = 0; + let dfs_traversal = std::iter::once(dfs_root) + .chain(graph.dfs(dfs_root).map(|(_, v)| v)) + .collect::>(); + + assert_eq!(dfs_traversal, vec![0, 3, 4, 2, 1]); + } + + #[test] + fn test_dfs_space_complexity() { + let num_v = 20; + let mut graph = Graph::new(num_v, 0); + for i in 0..num_v { + for j in 0..num_v { + graph.add_undirected_edge(i, j); + } + } + + let dfs_root = 7; + let mut dfs_search = graph.dfs(dfs_root); + let mut dfs_check = vec![dfs_root]; + for _ in 1..num_v { + dfs_check.push(dfs_search.next().unwrap().1); + assert!(dfs_search.stack.len() <= num_v + 1); + } + + dfs_check.sort(); + dfs_check.dedup(); + assert_eq!(0, dfs_check[0]); + assert_eq!(num_v, dfs_check.len()); + assert_eq!(num_v - 1, dfs_check[num_v - 1]); + } +} diff --git a/src/li_chao.rs b/src/li_chao.rs new file mode 100644 index 0000000..39cbb74 --- /dev/null +++ b/src/li_chao.rs @@ -0,0 +1,109 @@ +/// A structure for answering maximum queries on a set of linear functions. Supports two +/// operations: inserting a linear function and querying for maximum at a given point. +/// The queries can be done in any order, and we can do all the calculations using integers. +/// https://cp-algorithms.com/geometry/convex_hull_trick.html#li-chao-tree +/// Compared to the code in the above link, this implementation further improves the algorithm by +/// reducing the number of nodes to (right - left). This is done by removing the midpoint of a +/// segment from both children. Even better, this allows the index of a node to just be the +/// midpoint of the interval! +/// +/// Just like normal segment trees, this could be modified to a dynamic tree when the range is +/// huge, or if the queries are known in advance the x-coordinates can be compressed. +/// (it can also be made persistent!). +pub struct LiChaoTree { + left: i64, + right: i64, + lines: Vec<(i64, i64)>, +} + +impl LiChaoTree { + /// Creates a new tree, built to handle queries on the interval [left, right). + pub fn new(left: i64, right: i64) -> Self { + Self { + left, + right, + lines: vec![(0, i64::MIN); (right - left) as usize], + } + } + + /// Every node in the tree has the property that the line that maximizes its midpoint is found + /// either in the node or one of its ancestors. When we visit a node, we compute the winner at + /// the midpoint of the node. The winner is stored in the node. The loser can still possibly + /// beat the winner on some segment, either to the left or to the right of the current + /// midpoint, so we propagate it to that segment. This sequence ensures that the invariant is + /// kept. + fn max_with_impl(&mut self, mut m: i64, mut b: i64, l: i64, r: i64) { + if r <= l { + return; + } + let ix = ((r - self.left + l - self.left) / 2) as usize; + let mid = self.left + (ix as i64); + let (ref mut m_ix, ref mut b_ix) = self.lines[ix]; + if m * mid + b > *m_ix * mid + *b_ix { + std::mem::swap(&mut m, m_ix); + std::mem::swap(&mut b, b_ix); + } + if m < *m_ix { + self.max_with_impl(m, b, l, mid); + } else if m > *m_ix { + self.max_with_impl(m, b, mid + 1, r); + } + } + + /// Adds the line with slope m and intercept b. O(log N) complexity. + pub fn max_with(&mut self, m: i64, b: i64) { + self.max_with_impl(m, b, self.left, self.right); + } + + /// Because of the invariant established by add_line, we know that the best line for a given + /// point is stored in one of the ancestors of its node. So we accumulate the maximum answer as + /// we go back up the tree. + fn evaluate_impl(&self, x: i64, l: i64, r: i64) -> i64 { + if r == l { + return i64::MIN; + } + let ix = ((r - self.left + l - self.left) / 2) as usize; + let mid = ix as i64 + self.left; + let y = self.lines[ix].0 * x + self.lines[ix].1; + if x == mid { + y + } else if x < mid { + self.evaluate_impl(x, l, mid).max(y) + } else { + self.evaluate_impl(x, mid + 1, r).max(y) + } + } + + /// Finds the maximum mx+b among all lines in the structure. O(log N) complexity. + pub fn evaluate(&self, x: i64) -> i64 { + self.evaluate_impl(x, self.left, self.right) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_li_chao_tree() { + let lines = [(0, -3), (-1, 0), (1, -8), (-2, 1), (1, -4)]; + let xs = [0, 1, 2, 3, 4, 5]; + // results[i] consists of the expected y-coordinates after processing + // the first i+1 lines. + let results = [ + [-3, -3, -3, -3, -3, -3], + [0, -1, -2, -3, -3, -3], + [0, -1, -2, -3, -3, -3], + [1, -1, -2, -3, -3, -3], + [1, -1, -2, -1, 0, 1], + ]; + let mut li_chao = LiChaoTree::new(0, 6); + + assert_eq!(li_chao.evaluate(0), i64::MIN); + for (&(slope, intercept), expected) in lines.iter().zip(results.iter()) { + li_chao.max_with(slope, intercept); + let ys: Vec = xs.iter().map(|&x| li_chao.evaluate(x)).collect(); + assert_eq!(expected, &ys[..]); + } + } +} diff --git a/src/lib.rs b/src/lib.rs index b7fa839..ba482d1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,11 @@ //! Algorithms Cookbook in Rust. -pub mod arq_tree; + +pub mod caching; pub mod graph; +pub mod li_chao; pub mod math; +pub mod order; +pub mod range_query; +pub mod rng; pub mod scanner; pub mod string_proc; diff --git a/src/math.rs b/src/math.rs deleted file mode 100644 index 1671e50..0000000 --- a/src/math.rs +++ /dev/null @@ -1,65 +0,0 @@ -//! Number-theoretic utilities for contest problems. - -/// Modular exponentiation by repeated squaring: returns base^exp % m. -pub fn mod_pow(mut base: i64, mut exp: i64, m: i64) -> i64 { - let mut result = 1 % m; - while exp > 0 { - if exp % 2 == 1 { - result = (result * base) % m; - } - base = (base * base) % m; - exp /= 2; - } - result -} - -/// Finds (d, x, y) such that d = gcd(a, b) = ax + by. -pub fn extended_gcd(a: i64, b: i64) -> (i64, i64, i64) { - if b == 0 { - (a.abs(), a.signum(), 0) - } else { - let (d, x, y) = extended_gcd(b, a % b); - (d, y, x - y * (a / b)) - } -} - -/// Assuming a != 0, finds smallest y >= 0 such that ax + by = c. -pub fn canon_egcd(a: i64, b: i64, c: i64) -> Option<(i64, i64, i64)> { - let (d, _, yy) = extended_gcd(a, b); - if c % d == 0 { - let z = (a / d).abs(); - let y = (yy * (c / d) % z + z) % z; - let x = (c - b * y) / a; - Some((d, x, y)) - } else { - None - } -} - -#[cfg(test)] -mod test { - use super::*; - - #[test] - fn test_mod_inverse() { - let p = 1_000_000_007; - let base = 31; - - let base_inv = mod_pow(base, p - 2, p); - let identity = (base * base_inv) % p; - - assert_eq!(identity, 1); - } - - #[test] - fn test_egcd() { - let (a, b) = (14, 35); - - let (d, x, y) = extended_gcd(a, b); - assert_eq!(d, 7); - assert_eq!(a * x + b * y, d); - - assert_eq!(canon_egcd(a, b, d), Some((d, -2, 1))); - assert_eq!(canon_egcd(b, a, d), Some((d, -1, 3))); - } -} diff --git a/src/math/fft.rs b/src/math/fft.rs new file mode 100644 index 0000000..43d83dc --- /dev/null +++ b/src/math/fft.rs @@ -0,0 +1,236 @@ +//! The Fast Fourier Transform (FFT) and Number Theoretic Transform (NTT) +use super::num::{CommonField, Complex, PI}; +use std::ops::{Add, Div, Mul, Neg, Sub}; + +// We can delete this struct once f64::reverse_bits() stabilizes. +struct BitRevIterator { + a: usize, + n: usize, +} +impl BitRevIterator { + fn new(n: usize) -> Self { + assert!(n.is_power_of_two()); + Self { a: 2 * n - 1, n } + } +} +impl Iterator for BitRevIterator { + type Item = usize; + + fn next(&mut self) -> Option { + if self.a == 2 * self.n - 2 { + return None; + } + let mut mask = self.n; + while self.a & mask > 0 { + self.a ^= mask; + mask /= 2; + } + self.a |= mask; + Some(self.a / 2) + } +} + +#[allow(clippy::upper_case_acronyms)] +pub trait FFT: Sized + Copy { + type F: Sized + + Copy + + From + + Neg + + Add + + Div + + Mul + + Sub; + + const ZERO: Self; + + /// A primitive nth root of one raised to the powers 0, 1, 2, ..., n/2 - 1 + fn get_roots(n: usize, inverse: bool) -> Vec; + /// 1 for forward transform, 1/n for inverse transform + fn get_factor(n: usize, inverse: bool) -> Self::F; + /// The inverse of Self::F::from() + fn extract(f: Self::F) -> Self; +} + +impl FFT for f64 { + type F = Complex; + + const ZERO: f64 = 0.0; + + fn get_roots(n: usize, inverse: bool) -> Vec { + let step = if inverse { -2.0 } else { 2.0 } * PI / n as f64; + (0..n / 2) + .map(|i| Complex::from_polar(1.0, step * i as f64)) + .collect() + } + + fn get_factor(n: usize, inverse: bool) -> Self::F { + Self::F::from(if inverse { (n as f64).recip() } else { 1.0 }) + } + + fn extract(f: Self::F) -> f64 { + f.real + } +} + +// NTT notes: see problem 30-6 in CLRS for details, keeping in mind that +// 2187 and 410692747 are inverses and 2^26th roots of 1 mod (7<<26)+1 +// 15311432 and 469870224 are inverses and 2^23rd roots of 1 mod (119<<23)+1 +// 440564289 and 1713844692 are inverses and 2^27th roots of 1 mod (15<<27)+1 +// 125 and 2267742733 are inverses and 2^30th roots of 1 mod (3<<30)+1 +impl FFT for i64 { + type F = CommonField; + + const ZERO: Self = 0; + + fn get_roots(n: usize, inverse: bool) -> Vec { + assert!(n <= 1 << 23); + let mut prim_root = Self::F::from(15_311_432); + if inverse { + prim_root = prim_root.recip(); + } + for _ in (0..).take_while(|&i| n < 1 << (23 - i)) { + prim_root = prim_root * prim_root; + } + + let mut roots = Vec::with_capacity(n / 2); + let mut root = Self::F::from(1); + for _ in 0..roots.capacity() { + roots.push(root); + root = root * prim_root; + } + roots + } + + fn get_factor(n: usize, inverse: bool) -> Self::F { + Self::F::from(if inverse { n as Self } else { 1 }).recip() + } + + fn extract(f: Self::F) -> Self { + f.val + } +} + +/// Computes the discrete fourier transform of v, whose length is a power of 2. +/// Forward transform: polynomial coefficients -> evaluate at roots of unity +/// Inverse transform: values at roots of unity -> interpolated coefficients +pub fn fft(v: &[T::F], inverse: bool) -> Vec { + let n = v.len(); + assert!(n.is_power_of_two()); + + let factor = T::get_factor(n, inverse); + let roots_of_unity = T::get_roots(n, inverse); + let mut dft = BitRevIterator::new(n) + .map(|i| v[i] * factor) + .collect::>(); + + for m in (0..).map(|s| 1 << s).take_while(|&m| m < n) { + for k in (0..n).step_by(2 * m) { + for j in 0..m { + let u = dft[k + j]; + let t = dft[k + j + m] * roots_of_unity[n / 2 / m * j]; + dft[k + j] = u + t; + dft[k + j + m] = u - t; + } + } + } + dft +} + +/// From a slice of reals (f64 or i64), computes DFT of size at least desired_len +pub fn dft_from_reals(v: &[T], desired_len: usize) -> Vec { + assert!(v.len() <= desired_len); + + let complex_v = v + .iter() + .cloned() + .chain(std::iter::repeat(T::ZERO)) + .take(desired_len.next_power_of_two()) + .map(T::F::from) + .collect::>(); + fft::(&complex_v, false) +} + +/// The inverse of dft_from_reals() +pub fn idft_to_reals(dft_v: &[T::F], desired_len: usize) -> Vec { + assert!(dft_v.len() >= desired_len); + + let complex_v = fft::(dft_v, true); + complex_v + .into_iter() + .take(desired_len) + .map(T::extract) + .collect() +} + +/// Given two polynomials (vectors) sum_i a[i] x^i and sum_i b[i] x^i, +/// computes their product (convolution) c[k] = sum_(i+j=k) a[i]*b[j]. +/// Uses complex FFT if inputs are f64, or modular NTT if inputs are i64. +pub fn convolution(a: &[T], b: &[T]) -> Vec { + let len_c = a.len() + b.len() - 1; + let dft_a = dft_from_reals(a, len_c).into_iter(); + let dft_b = dft_from_reals(b, len_c).into_iter(); + let dft_c = dft_a.zip(dft_b).map(|(a, b)| a * b).collect::>(); + idft_to_reals(&dft_c, len_c) +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_complex_dft() { + let v = vec![7.0, 1.0, 1.0]; + let dft_v = dft_from_reals(&v, v.len()); + let new_v: Vec = idft_to_reals(&dft_v, v.len()); + + let six = Complex::from(6.0); + let seven = Complex::from(7.0); + let nine = Complex::from(9.0); + let i = Complex::new(0.0, 1.0); + + assert_eq!(dft_v, vec![nine, six + i, seven, six - i]); + assert_eq!(new_v, v); + } + + #[test] + fn test_modular_dft() { + let v = vec![7, 1, 1]; + let dft_v = dft_from_reals(&v, v.len()); + let new_v: Vec = idft_to_reals(&dft_v, v.len()); + + let seven = CommonField::from(7); + let one = CommonField::from(1); + let prim = CommonField::from(15_311_432).pow(1 << 21); + let prim2 = prim * prim; + + let eval0 = seven + one + one; + let eval1 = seven + prim + prim2; + let eval2 = seven + prim2 + one; + let eval3 = seven + prim.recip() + prim2; + + assert_eq!(dft_v, vec![eval0, eval1, eval2, eval3]); + assert_eq!(new_v, v); + } + + #[test] + fn test_complex_convolution() { + let x = vec![7.0, 1.0, 1.0]; + let y = vec![2.0, 4.0]; + let z = convolution(&x, &y); + let m = convolution(&vec![999.0], &vec![1e6]); + + assert_eq!(z, vec![14.0, 30.0, 6.0, 4.0]); + assert_eq!(m, vec![999e6]); + } + + #[test] + fn test_modular_convolution() { + let x = vec![7, 1, 1]; + let y = vec![2, 4]; + let z = convolution(&x, &y); + let m = convolution(&vec![999], &vec![1_000_000]); + + assert_eq!(z, vec![14, 30, 6, 4]); + assert_eq!(m, vec![999_000_000 - super::super::num::COMMON_PRIME]); + } +} diff --git a/src/math/mod.rs b/src/math/mod.rs new file mode 100644 index 0000000..c27e7af --- /dev/null +++ b/src/math/mod.rs @@ -0,0 +1,176 @@ +//! Number-theoretic utilities for contest problems. +pub mod fft; +pub mod num; + +/// Finds (d, coef_a, coef_b) such that d = gcd(a, b) = a * coef_a + b * coef_b. +pub fn extended_gcd(a: i64, b: i64) -> (i64, i64, i64) { + if b == 0 { + (a.abs(), a.signum(), 0) + } else { + let (d, coef_b, coef_a) = extended_gcd(b, a % b); + (d, coef_a, coef_b - coef_a * (a / b)) + } +} + +/// Assuming a != 0, finds smallest coef_b >= 0 such that a * coef_a + b * coef_b = c. +/// +/// # Panics +/// +/// Panics if a == 0. +pub fn canon_egcd(a: i64, b: i64, c: i64) -> Option<(i64, i64, i64)> { + let (d, _, coef_b_init) = extended_gcd(a, b); + if c % d == 0 { + let a_d = (a / d).abs(); + let coef_b = (coef_b_init * (c / d) % a_d + a_d) % a_d; + let coef_a = (c - b * coef_b) / a; + Some((d, coef_a, coef_b)) + } else { + None + } +} + +// TODO: deduplicate modular arithmetic code with num::Field +fn pos_mod(n: i64, m: i64) -> i64 { + if n < 0 { n + m } else { n } +} +fn mod_mul(a: i64, b: i64, m: i64) -> i64 { + pos_mod((a as i128 * b as i128 % m as i128) as i64, m) +} +fn mod_exp(mut base: i64, mut exp: u64, m: i64) -> i64 { + assert!(m >= 1); + let mut ans = 1 % m; + base %= m; + while exp > 0 { + if exp % 2 == 1 { + ans = mod_mul(ans, base, m); + } + base = mod_mul(base, base, m); + exp /= 2; + } + pos_mod(ans, m) +} + +fn is_strong_probable_prime(n: i64, exp: u64, r: i64, a: i64) -> bool { + let mut x = mod_exp(a, exp, n); + if x == 1 || x == n - 1 { + return true; + } + for _ in 1..r { + x = mod_mul(x, x, n); + if x == n - 1 { + return true; + } + } + false +} + +/// Assuming x >= 0, returns whether x is prime +pub fn is_prime(n: i64) -> bool { + const BASES: [i64; 12] = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37]; + assert!(n >= 0); + match n { + 0 | 1 => false, + 2 | 3 => true, + _ if n % 2 == 0 => false, + _ => { + let r = (n - 1).trailing_zeros() as i64; + let exp = (n - 1) as u64 >> r; + BASES + .iter() + .all(|&base| base > n - 2 || is_strong_probable_prime(n, exp, r, base)) + } + } +} + +fn pollard_rho(n: i64) -> i64 { + for a in 1..n { + let f = |x| pos_mod(mod_mul(x, x, n) + a, n); + let mut x = 2; + let mut y = 2; + loop { + x = f(x); + y = f(f(y)); + let div = num::fast_gcd(x - y, n); + if div == n { + break; + } else if div > 1 { + return div; + } + } + } + panic!("No divisor found!"); +} + +/// Assuming x >= 1, finds the prime factorization of n +/// TODO: pollard_rho needs randomization to ensure correctness in contest settings! +pub fn factorize(n: i64) -> Vec { + assert!(n >= 1); + let r = n.trailing_zeros() as usize; + let mut factors = vec![2; r]; + let mut stack = match n >> r { + 1 => vec![], + x => vec![x], + }; + while let Some(top) = stack.pop() { + if is_prime(top) { + factors.push(top); + } else { + let div = pollard_rho(top); + stack.push(div); + stack.push(top / div); + } + } + factors.sort_unstable(); + factors +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_egcd() { + let (a, b) = (14, 35); + + let (d, x, y) = extended_gcd(a, b); + assert_eq!(d, 7); + assert_eq!(a * x + b * y, d); + + assert_eq!(canon_egcd(a, b, d), Some((d, -2, 1))); + assert_eq!(canon_egcd(b, a, d), Some((d, -1, 3))); + } + + #[test] + fn test_modexp() { + let m = 1_000_000_007; + assert_eq!(mod_exp(0, 0, m), 1); + assert_eq!(mod_exp(0, 1, m), 0); + assert_eq!(mod_exp(0, 10, m), 0); + assert_eq!(mod_exp(123, 456, m), 565291922); + } + + #[test] + fn test_miller() { + assert_eq!(is_prime(2), true); + assert_eq!(is_prime(4), false); + assert_eq!(is_prime(6), false); + assert_eq!(is_prime(8), false); + assert_eq!(is_prime(269), true); + assert_eq!(is_prime(1000), false); + assert_eq!(is_prime(1_000_000_007), true); + assert_eq!(is_prime((1 << 61) - 1), true); + assert_eq!(is_prime(7156857700403137441), false); + } + + #[test] + fn test_pollard() { + assert_eq!(factorize(1), vec![]); + assert_eq!(factorize(2), vec![2]); + assert_eq!(factorize(4), vec![2, 2]); + assert_eq!(factorize(12), vec![2, 2, 3]); + assert_eq!( + factorize(7156857700403137441), + vec![11, 13, 17, 19, 29, 37, 41, 43, 61, 97, 109, 127] + ); + } +} diff --git a/src/math/num.rs b/src/math/num.rs new file mode 100644 index 0000000..ef37619 --- /dev/null +++ b/src/math/num.rs @@ -0,0 +1,466 @@ +//! Rational and Complex numbers, safe modular arithmetic, and linear algebra, +//! implemented minimally for contest use. +//! If you need more features, you might be interested in crates.io/crates/num +pub use std::f64::consts::PI; +use std::ops::{Add, Div, Index, IndexMut, Mul, Neg, Sub}; + +/// Fast iterative version of Euclid's GCD algorithm +pub fn fast_gcd(mut a: i64, mut b: i64) -> i64 { + while b != 0 { + a %= b; + std::mem::swap(&mut a, &mut b); + } + a.abs() +} + +/// Represents a fraction reduced to lowest terms +#[derive(Clone, Copy, Eq, PartialEq, Debug, Hash)] +pub struct Rational { + pub num: i64, + pub den: i64, +} +impl Rational { + pub fn new(num: i64, den: i64) -> Self { + let g = fast_gcd(num, den) * den.signum(); + Self { + num: num / g, + den: den / g, + } + } + pub fn abs(self) -> Self { + Self { + num: self.num.abs(), + den: self.den, + } + } + pub fn recip(self) -> Self { + let g = self.num.signum(); + Self { + num: self.den / g, + den: self.num / g, + } + } +} +impl From for Rational { + fn from(num: i64) -> Self { + Self { num, den: 1 } + } +} +impl Neg for Rational { + type Output = Self; + fn neg(self) -> Self { + Self { + num: -self.num, + den: self.den, + } + } +} +#[allow(clippy::suspicious_arithmetic_impl)] +impl Add for Rational { + type Output = Self; + fn add(self, other: Self) -> Self { + Self::new( + self.num * other.den + self.den * other.num, + self.den * other.den, + ) + } +} +#[allow(clippy::suspicious_arithmetic_impl)] +impl Sub for Rational { + type Output = Self; + fn sub(self, other: Self) -> Self { + Self::new( + self.num * other.den - self.den * other.num, + self.den * other.den, + ) + } +} +impl Mul for Rational { + type Output = Self; + fn mul(self, other: Self) -> Self { + Self::new(self.num * other.num, self.den * other.den) + } +} +#[allow(clippy::suspicious_arithmetic_impl)] +impl Div for Rational { + type Output = Self; + fn div(self, other: Self) -> Self { + self * other.recip() + } +} +impl Ord for Rational { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + (self.num * other.den).cmp(&(self.den * other.num)) + } +} +impl PartialOrd for Rational { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +/// Represents a complex number using floating-point arithmetic +#[derive(Clone, Copy, PartialEq, Debug)] +pub struct Complex { + pub real: f64, + pub imag: f64, +} +impl Complex { + pub fn new(real: f64, imag: f64) -> Self { + Self { real, imag } + } + pub fn from_polar(r: f64, th: f64) -> Self { + Self::new(r * th.cos(), r * th.sin()) + } + pub fn abs_square(self) -> f64 { + self.real * self.real + self.imag * self.imag + } + pub fn argument(self) -> f64 { + self.imag.atan2(self.real) + } + pub fn conjugate(self) -> Self { + Self::new(self.real, -self.imag) + } + pub fn recip(self) -> Self { + let denom = self.abs_square(); + Self::new(self.real / denom, -self.imag / denom) + } +} +impl From for Complex { + fn from(real: f64) -> Self { + Self::new(real, 0.0) + } +} +impl Neg for Complex { + type Output = Self; + fn neg(self) -> Self { + Self::new(-self.real, -self.imag) + } +} +impl Add for Complex { + type Output = Self; + fn add(self, other: Self) -> Self { + Self::new(self.real + other.real, self.imag + other.imag) + } +} +impl Sub for Complex { + type Output = Self; + fn sub(self, other: Self) -> Self { + Self::new(self.real - other.real, self.imag - other.imag) + } +} +impl Mul for Complex { + type Output = Self; + fn mul(self, other: Self) -> Self { + let real = self.real * other.real - self.imag * other.imag; + let imag = self.imag * other.real + self.real * other.imag; + Self::new(real, imag) + } +} +#[allow(clippy::suspicious_arithmetic_impl)] +impl Div for Complex { + type Output = Self; + fn div(self, other: Self) -> Self { + self * other.recip() + } +} + +/// Represents an element of the finite (Galois) field of prime order M, where +/// 1 <= M < 2^31.5. If M is not prime, ring operations are still valid +/// but recip() and division are not. Note that the latter operations are also +/// the slowest, so precompute any inverses that you intend to use frequently. +#[derive(Clone, Copy, Eq, PartialEq, Debug, Hash)] +pub struct Modulo { + pub val: i64, +} +impl Modulo { + /// Computes self^n in O(log n) time + pub fn pow(mut self, mut n: u64) -> Self { + let mut result = Self::from_small(1); + while n > 0 { + if n % 2 == 1 { + result = result * self; + } + self = self * self; + n /= 2; + } + result + } + /// Computes inverses of 1 to n in O(n) time + pub fn vec_of_recips(n: i64) -> Vec { + let mut recips = vec![Self::from(0), Self::from(1)]; + for i in 2..=n { + let (md, dv) = (M % i, M / i); + recips.push(recips[md as usize] * Self::from_small(-dv)); + } + recips + } + /// Computes self^-1 in O(log M) time + pub fn recip(self) -> Self { + self.pow(M as u64 - 2) + } + /// Avoids the % operation but requires -M <= x < M + fn from_small(s: i64) -> Self { + let val = if s < 0 { s + M } else { s }; + Self { val } + } +} +impl From for Modulo { + fn from(val: i64) -> Self { + // Self { val: val.rem_euclid(M) } + Self::from_small(val % M) + } +} +impl Neg for Modulo { + type Output = Self; + fn neg(self) -> Self { + Self::from_small(-self.val) + } +} +impl Add for Modulo { + type Output = Self; + fn add(self, other: Self) -> Self { + Self::from_small(self.val + other.val - M) + } +} +impl Sub for Modulo { + type Output = Self; + fn sub(self, other: Self) -> Self { + Self::from_small(self.val - other.val) + } +} +impl Mul for Modulo { + type Output = Self; + fn mul(self, other: Self) -> Self { + Self::from(self.val * other.val) + } +} +#[allow(clippy::suspicious_arithmetic_impl)] +impl Div for Modulo { + type Output = Self; + fn div(self, other: Self) -> Self { + self * other.recip() + } +} + +/// Prime modulus that's commonly used in programming competitions +pub const COMMON_PRIME: i64 = 998_244_353; // 2^23 * 7 * 17 + 1; +pub type CommonField = Modulo; + +#[derive(Clone, PartialEq, Debug)] +pub struct Matrix { + cols: usize, + inner: Box<[f64]>, +} +impl Matrix { + pub fn zero(rows: usize, cols: usize) -> Self { + let inner = vec![0.0; rows * cols].into_boxed_slice(); + Self { cols, inner } + } + pub fn one(cols: usize) -> Self { + let mut matrix = Self::zero(cols, cols); + for i in 0..cols { + matrix[i][i] = 1.0; + } + matrix + } + pub fn vector(vec: &[f64], as_row: bool) -> Self { + let cols = if as_row { vec.len() } else { 1 }; + let inner = vec.to_vec().into_boxed_slice(); + Self { cols, inner } + } + pub fn pow(&self, mut n: u64) -> Self { + let mut base = self.clone(); + let mut result = Self::one(self.cols); + while n > 0 { + if n % 2 == 1 { + result = &result * &base; + } + base = &base * &base; + n /= 2; + } + result + } + pub fn rows(&self) -> usize { + self.inner.len() / self.cols + } + pub fn transpose(&self) -> Self { + let mut matrix = Matrix::zero(self.cols, self.rows()); + for i in 0..self.rows() { + for j in 0..self.cols { + matrix[j][i] = self[i][j]; + } + } + matrix + } + pub fn recip(&self) -> Self { + unimplemented!(); + } +} +impl Index for Matrix { + type Output = [f64]; + fn index(&self, row: usize) -> &Self::Output { + let start = self.cols * row; + &self.inner[start..start + self.cols] + } +} +impl IndexMut for Matrix { + fn index_mut(&mut self, row: usize) -> &mut Self::Output { + let start = self.cols * row; + &mut self.inner[start..start + self.cols] + } +} +impl Neg for &Matrix { + type Output = Matrix; + fn neg(self) -> Matrix { + let inner = self.inner.iter().map(|&v| -v).collect(); + Matrix { + cols: self.cols, + inner, + } + } +} +impl Add for &Matrix { + type Output = Matrix; + fn add(self, other: Self) -> Matrix { + let self_iter = self.inner.iter(); + let inner = self_iter + .zip(other.inner.iter()) + .map(|(&u, &v)| u + v) + .collect(); + Matrix { + cols: self.cols, + inner, + } + } +} +impl Sub for &Matrix { + type Output = Matrix; + fn sub(self, other: Self) -> Matrix { + let self_iter = self.inner.iter(); + let inner = self_iter + .zip(other.inner.iter()) + .map(|(&u, &v)| u - v) + .collect(); + Matrix { + cols: self.cols, + inner, + } + } +} +impl Mul for &Matrix { + type Output = Matrix; + fn mul(self, scalar: f64) -> Matrix { + let inner = self.inner.iter().map(|&v| v * scalar).collect(); + Matrix { + cols: self.cols, + inner, + } + } +} +impl Mul for &Matrix { + type Output = Matrix; + fn mul(self, other: Self) -> Matrix { + assert_eq!(self.cols, other.rows()); + let mut matrix = Matrix::zero(self.rows(), other.cols); + for i in 0..self.rows() { + for k in 0..self.cols { + for j in 0..other.cols { + matrix[i][j] += self[i][k] * other[k][j]; + } + } + } + matrix + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_rational() { + let three = Rational::from(3); + let six = Rational::from(6); + let three_and_half = three + three / six; + + assert_eq!(three_and_half.num, 7); + assert_eq!(three_and_half.den, 2); + assert_eq!(three_and_half, Rational::new(-35, -10)); + assert!(three_and_half > Rational::from(3)); + assert!(three_and_half < Rational::from(4)); + + let minus_three_and_half = six - three_and_half + three / (-three / six); + let zero = three_and_half + minus_three_and_half; + + assert_eq!(minus_three_and_half.num, -7); + assert_eq!(minus_three_and_half.den, 2); + assert_eq!(three_and_half, -minus_three_and_half); + assert_eq!(zero.num, 0); + assert_eq!(zero.den, 1); + } + + #[test] + fn test_complex() { + let four = Complex::new(4.0, 0.0); + let two_i = Complex::new(0.0, 2.0); + + assert_eq!(four / two_i, -two_i); + assert_eq!(two_i * -two_i, four); + assert_eq!(two_i - two_i, Complex::from(0.0)); + assert_eq!(four.abs_square(), 16.0); + assert_eq!(two_i.abs_square(), 4.0); + assert_eq!((-four).argument(), -PI); + assert_eq!((-two_i).argument(), -PI / 2.0); + assert_eq!(four.argument(), 0.0); + assert_eq!(two_i.argument(), PI / 2.0); + } + + #[test] + fn test_field() { + let base = CommonField::from(1234); + let zero = base - base; + let one = base.recip() * base; + let two = CommonField::from(2 - 5 * COMMON_PRIME); + + assert_eq!(zero.val, 0); + assert_eq!(one.val, 1); + assert_eq!(one + one, two); + assert_eq!(one / base * (base * base) - base / one, zero); + } + + #[test] + fn test_vec_of_recips() { + let recips = CommonField::vec_of_recips(20); + + assert_eq!(recips.len(), 21); + for i in 1..recips.len() { + assert_eq!(recips[i], CommonField::from(i as i64).recip()); + } + } + + #[test] + fn test_linalg() { + let zero = Matrix::zero(2, 2); + let one = Matrix::one(2); + let rotate_90 = Matrix { + cols: 2, + inner: Box::new([0.0, -1.0, 1.0, 0.0]), + }; + let x_vec = Matrix::vector(&[1.0, 0.0], false); + let y_vec = Matrix::vector(&[0.0, 1.0], false); + let x_dot_x = &x_vec.transpose() * &x_vec; + let x_dot_y = &x_vec.transpose() * &y_vec; + + assert_eq!(x_dot_x, Matrix::one(1)); + assert_eq!(x_dot_x[0][0], 1.0); + assert_eq!(x_dot_y, Matrix::zero(1, 1)); + assert_eq!(x_dot_y[0][0], 0.0); + assert_eq!(&one - &one, zero); + assert_eq!(&one * 0.0, zero); + assert_eq!(&rotate_90 * &rotate_90, -&one); + assert_eq!(&rotate_90 * &x_vec, y_vec); + assert_eq!(&rotate_90 * &y_vec, -&x_vec); + assert_eq!(&rotate_90 * &(&x_vec + &y_vec), &y_vec - &x_vec); + } +} diff --git a/src/order.rs b/src/order.rs new file mode 100644 index 0000000..6ba82d9 --- /dev/null +++ b/src/order.rs @@ -0,0 +1,233 @@ +//! Ordering algorithms. + +/// A comparator on partially ordered elements, that panics if they are incomparable +/// +/// # Example +/// +/// ``` +/// use contest_algorithms::order::asserting_cmp; +/// let mut vec = vec![4.5, -1.7, 1.2]; +/// vec.sort_unstable_by(asserting_cmp); +/// assert_eq!(vec, vec![-1.7, 1.2, 4.5]); +/// ``` +pub fn asserting_cmp(a: &T, b: &T) -> std::cmp::Ordering { + a.partial_cmp(b).expect("Comparing incomparable elements") +} + +/// Assuming slice is sorted and totally ordered, returns the minimum i for which +/// slice[i] >= key, or slice.len() if no such i exists +pub fn slice_lower_bound(slice: &[T], key: &T) -> usize { + slice + .binary_search_by(|x| asserting_cmp(x, key).then(std::cmp::Ordering::Greater)) + .unwrap_err() +} + +/// Assuming slice is sorted and totally ordered, returns the minimum i for which +/// slice[i] > key, or slice.len() if no such i exists +pub fn slice_upper_bound(slice: &[T], key: &T) -> usize { + slice + .binary_search_by(|x| asserting_cmp(x, key).then(std::cmp::Ordering::Less)) + .unwrap_err() +} + +/// Stably merges two sorted and totally ordered collections into one +pub fn merge_sorted( + i1: impl IntoIterator, + i2: impl IntoIterator, +) -> Vec { + let mut i1 = i1.into_iter().peekable(); + let mut i2 = i2.into_iter().peekable(); + let mut merged = Vec::with_capacity(i1.size_hint().0 + i2.size_hint().0); + while let (Some(a), Some(b)) = (i1.peek(), i2.peek()) { + merged.push(if a <= b { i1.next() } else { i2.next() }.unwrap()); + } + merged.extend(i1.chain(i2)); + merged +} + +/// A stable sort +pub fn merge_sort(mut v: Vec) -> Vec { + if v.len() < 2 { + v + } else { + let v2 = v.split_off(v.len() / 2); + merge_sorted(merge_sort(v), merge_sort(v2)) + } +} + +/// A simple data structure for coordinate compression +pub struct SparseIndex { + coords: Vec, +} + +impl SparseIndex { + /// Builds an index, given the full set of coordinates to compress. + pub fn new(mut coords: Vec) -> Self { + coords.sort_unstable(); + coords.dedup(); + Self { coords } + } + + /// Returns Ok(i) if the coordinate q appears at index i + /// Returns Err(i) if q appears between indices i-1 and i + pub fn compress(&self, q: i64) -> Result { + self.coords.binary_search(&q) + } +} + +/// Represents a maximum (upper envelope) of a collection of linear functions of one +/// variable, evaluated using an online version of the convex hull trick. +/// It combines the offline algorithm with square root decomposition, resulting in an +/// asymptotically suboptimal but simple algorithm with good amortized performance: +/// N inserts interleaved with Q queries yields O(N sqrt Q + Q log N) time complexity +/// in general, or O((N + Q) log N) if all queries come after all inserts. +// Proof: the Q log N term comes from calls to slice_lower_bound(). As for the N sqrt Q, +// note that between successive times when the hull is rebuilt, O(N) work is done, +// and the running totals of insertions and queries satisfy del_N (del_Q + 1) > N. +// Now, either del_Q >= sqrt Q, or else del_Q <= 2 sqrt Q - 1 +// => del_N > N / (2 sqrt Q). +// Since del(N sqrt Q) >= max(N del(sqrt Q), del_N sqrt Q) +// >= max(N del_Q / (2 sqrt Q), del_N sqrt Q), +// we conclude that del(N sqrt Q) >= N / 2. +#[derive(Default)] +pub struct PiecewiseLinearConvexFn { + recent_lines: Vec<(f64, f64)>, + sorted_lines: Vec<(f64, f64)>, + intersections: Vec, + amortized_work: usize, +} + +impl PiecewiseLinearConvexFn { + /// Replaces the represented function with the maximum of itself and a provided line + pub fn max_with(&mut self, new_m: f64, new_b: f64) { + self.recent_lines.push((new_m, new_b)); + } + + /// Similar to max_with but requires that (new_m, new_b) be the largest pair so far + fn max_with_sorted(&mut self, new_m: f64, new_b: f64) { + while let Some(&(last_m, last_b)) = self.sorted_lines.last() { + // If slopes are equal, get rid of the old line as its intercept is lower + if (new_m - last_m).abs() > 1e-9 { + let intersect = (new_b - last_b) / (last_m - new_m); + if self.intersections.last() < Some(&intersect) { + self.intersections.push(intersect); + break; + } + } + self.intersections.pop(); + self.sorted_lines.pop(); + } + self.sorted_lines.push((new_m, new_b)); + } + + /// Evaluates the function at x + fn eval_unoptimized(&self, x: f64) -> f64 { + let idx = slice_lower_bound(&self.intersections, &x); + self.recent_lines + .iter() + .chain(self.sorted_lines.get(idx)) + .map(|&(m, b)| m * x + b) + .max_by(asserting_cmp) + .unwrap_or(-1e18) + } + + /// Evaluates the function at x with good amortized runtime + pub fn evaluate(&mut self, x: f64) -> f64 { + self.amortized_work += self.recent_lines.len(); + if self.amortized_work > self.sorted_lines.len() { + self.amortized_work = 0; + self.recent_lines.sort_unstable_by(asserting_cmp); + self.intersections.clear(); + let all_lines = merge_sorted(self.recent_lines.drain(..), self.sorted_lines.drain(..)); + for (new_m, new_b) in all_lines { + self.max_with_sorted(new_m, new_b); + } + } + self.eval_unoptimized(x) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_bounds() { + let mut vals = vec![16, 45, 45, 45, 82]; + + assert_eq!(slice_upper_bound(&vals, &44), 1); + assert_eq!(slice_lower_bound(&vals, &45), 1); + assert_eq!(slice_upper_bound(&vals, &45), 4); + assert_eq!(slice_lower_bound(&vals, &46), 4); + + vals.dedup(); + for (i, q) in vals.iter().enumerate() { + assert_eq!(slice_lower_bound(&vals, q), i); + assert_eq!(slice_upper_bound(&vals, q), i + 1); + } + } + + #[test] + fn test_merge_sorted() { + let vals1 = vec![16, 45, 45, 82]; + let vals2 = vec![-20, 40, 45, 50]; + let vals_merged = vec![-20, 16, 40, 45, 45, 45, 50, 82]; + + assert_eq!(merge_sorted(None, Some(42)), vec![42]); + assert_eq!(merge_sorted(vals1.iter().cloned(), None), vals1); + assert_eq!(merge_sorted(vals1, vals2), vals_merged); + } + + #[test] + fn test_merge_sort() { + let unsorted = vec![8, -5, 1, 4, -3, 4]; + let sorted = vec![-5, -3, 1, 4, 4, 8]; + + assert_eq!(merge_sort(unsorted), sorted); + assert_eq!(merge_sort(sorted.clone()), sorted); + } + + #[test] + fn test_coord_compress() { + let mut coords = vec![16, 99, 45, 18]; + let index = SparseIndex::new(coords.clone()); + + coords.sort_unstable(); + for (i, q) in coords.into_iter().enumerate() { + assert_eq!(index.compress(q - 1), Err(i)); + assert_eq!(index.compress(q), Ok(i)); + assert_eq!(index.compress(q + 1), Err(i + 1)); + } + } + + #[test] + fn test_range_compress() { + let queries = vec![(0, 10), (10, 19), (20, 29)]; + let coords = queries.iter().flat_map(|&(i, j)| vec![i, j + 1]).collect(); + let index = SparseIndex::new(coords); + + assert_eq!(index.coords, vec![0, 10, 11, 20, 30]); + } + + #[test] + fn test_convex_hull_trick() { + let lines = [(0, -3), (-1, 0), (1, -8), (-2, 1), (1, -4)]; + let xs = [0, 1, 2, 3, 4, 5]; + // results[i] consists of the expected y-coordinates after processing + // the first i+1 lines. + let results = [ + [-3, -3, -3, -3, -3, -3], + [0, -1, -2, -3, -3, -3], + [0, -1, -2, -3, -3, -3], + [1, -1, -2, -3, -3, -3], + [1, -1, -2, -1, 0, 1], + ]; + let mut func = PiecewiseLinearConvexFn::default(); + assert_eq!(func.evaluate(0.0), -1e18); + for (&(slope, intercept), expected) in lines.iter().zip(results.iter()) { + func.max_with(slope as f64, intercept as f64); + let ys: Vec = xs.iter().map(|&x| func.evaluate(x as f64) as i64).collect(); + assert_eq!(expected, &ys[..]); + } + } +} diff --git a/src/range_query/README.md b/src/range_query/README.md new file mode 100644 index 0000000..aeb953a --- /dev/null +++ b/src/range_query/README.md @@ -0,0 +1,3 @@ +# Associative Range Query (ARQ) and Mo's Algorithm + +For more information on Associative Range Query, you may research "segment trees" in the programming contest literature. My implementation is more general than usual; for more information on it, please see my [blog post on Codeforces](https://codeforces.com/blog/entry/68419). diff --git a/src/range_query/dynamic_arq.rs b/src/range_query/dynamic_arq.rs new file mode 100644 index 0000000..6450468 --- /dev/null +++ b/src/range_query/dynamic_arq.rs @@ -0,0 +1,185 @@ +//! Associative Range Query Tree with dynamic allocation, supporting sparse +//! initialization and persistence +use super::ArqSpec; + +pub struct DynamicArqNode { + val: T::S, + app: Option, + down: (usize, usize), +} + +// TODO: in a future Rust version, this might be replaced by a #[derive(Clone)] +impl Clone for DynamicArqNode { + fn clone(&self) -> Self { + Self { + val: self.val.clone(), + app: self.app.clone(), + down: self.down, + } + } +} + +impl Default for DynamicArqNode { + fn default() -> Self { + Self { + val: T::identity(), + app: None, + down: (usize::MAX, usize::MAX), + } + } +} + +impl DynamicArqNode { + fn apply(&mut self, f: &T::F, size: i64) { + self.val = T::apply(f, &self.val, size); + if size > 1 { + let h = match self.app { + Some(ref g) => T::compose(f, g), + None => f.clone(), + }; + self.app = Some(h); + } + } +} + +pub type ArqView = (usize, i64); + +/// A dynamic, and optionally persistent, associative range query data structure. +pub struct DynamicArq { + nodes: Vec>, + is_persistent: bool, +} + +impl DynamicArq { + /// Initializes the data structure without creating any nodes. + pub fn new(is_persistent: bool) -> Self { + Self { + nodes: vec![], + is_persistent, + } + } + + /// Lazily builds a tree initialized to the identity. + pub fn build_from_identity(&mut self, size: i64) -> ArqView { + self.nodes.push(DynamicArqNode::default()); + (self.nodes.len() - 1, size) + } + + /// Builds a tree whose leaves are set to a given non-empty slice. + pub fn build_from_slice(&mut self, init_val: &[T::S]) -> ArqView { + if init_val.len() == 1 { + let root = DynamicArqNode { + val: init_val[0].clone(), + ..Default::default() + }; + self.nodes.push(root); + (self.nodes.len() - 1, 1) + } else { + let ls = init_val.len() / 2; + let (l_init, r_init) = init_val.split_at(ls); + let l_view = self.build_from_slice(l_init); + let r_view = self.build_from_slice(r_init); + self.merge_equal_sized(l_view, r_view) + } + } + + /// Merges two balanced subtrees into a single tree with a 0-indexed view. + pub fn merge_equal_sized(&mut self, (lp, ls): ArqView, (rp, rs): ArqView) -> ArqView { + assert!(ls == rs || ls + 1 == rs); + let p = self.nodes.len(); + let root = DynamicArqNode { + down: (lp, rp), + ..Default::default() + }; + self.nodes.push(root); + self.pull(p); + (p, ls + rs) + } + + pub fn push(&mut self, (p, s): ArqView) -> (ArqView, ArqView) { + if self.nodes[p].down.0 == usize::MAX { + self.nodes.push(DynamicArqNode::default()); + self.nodes.push(DynamicArqNode::default()); + self.nodes[p].down = (self.nodes.len() - 2, self.nodes.len() - 1) + }; + let (lp, rp) = self.nodes[p].down; + let ls = s / 2; + if let Some(ref f) = self.nodes[p].app.take() { + self.nodes[lp].apply(f, ls); + self.nodes[rp].apply(f, s - ls); + } + ((lp, ls), (rp, s - ls)) + } + + pub fn pull(&mut self, p: usize) { + let (lp, rp) = self.nodes[p].down; + let left_val = &self.nodes[lp].val; + let right_val = &self.nodes[rp].val; + self.nodes[p].val = T::op(left_val, right_val); + } + + fn clone_node(&mut self, p_orig: usize) -> usize { + if self.is_persistent { + let node = self.nodes[p_orig].clone(); + self.nodes.push(node); + self.nodes.len() - 1 + } else { + p_orig + } + } + + /// Applies the endomorphism f to all entries from l to r, inclusive. + /// If l == r, the updates are eager. Otherwise, they are lazy. + pub fn update(&mut self, view: ArqView, l: i64, r: i64, f: &T::F) -> ArqView { + let (p_orig, s) = view; + if r < 0 || s - 1 < l { + view + } else if l <= 0 && s - 1 <= r { + let p_clone = self.clone_node(p_orig); + self.nodes[p_clone].apply(f, s); + (p_clone, s) + } else { + let (l_view, r_view) = self.push(view); + let ls = l_view.1; + let p_clone = self.clone_node(p_orig); + let lp_clone = self.update(l_view, l, r, f).0; + let rp_clone = self.update(r_view, l - ls, r - ls, f).0; + self.nodes[p_clone].down = (lp_clone, rp_clone); + self.pull(p_clone); + (p_clone, s) + } + } + + /// Returns the aggregate range query on all entries from l to r, inclusive. + pub fn query(&mut self, view: ArqView, l: i64, r: i64) -> T::S { + let (p, s) = view; + if r < 0 || s - 1 < l { + T::identity() + } else if l <= 0 && s - 1 <= r { + self.nodes[p].val.clone() + } else { + let (l_view, r_view) = self.push(view); + let ls = l_view.1; + let l_agg = self.query(l_view, l, r); + let r_agg = self.query(r_view, l - ls, r - ls); + T::op(&l_agg, &r_agg) + } + } +} + +/// An example of binary search to find the first position whose element is negative. +/// The DynamicArq version works on trees of any size, not necessarily a power of two. +pub fn first_negative(arq: &mut DynamicArq, view: ArqView) -> Option { + let (p, s) = view; + if s == 1 { + Some(0).filter(|_| arq.nodes[p].val < 0) + } else { + let (l_view, r_view) = arq.push(view); + let (lp, ls) = l_view; + if arq.nodes[lp].val < 0 { + first_negative(arq, l_view) + } else { + first_negative(arq, r_view).map(|x| ls + x) + } + } +} diff --git a/src/range_query/mod.rs b/src/range_query/mod.rs new file mode 100644 index 0000000..332fb59 --- /dev/null +++ b/src/range_query/mod.rs @@ -0,0 +1,146 @@ +pub mod dynamic_arq; +pub mod specs; +pub mod sqrt_decomp; +pub mod static_arq; +pub use dynamic_arq::{ArqView, DynamicArq}; +pub use specs::ArqSpec; +pub use static_arq::StaticArq; + +#[cfg(test)] +mod test { + use super::specs::*; + use super::*; + + #[test] + fn test_rmq() { + let mut arq = StaticArq::::new(&[0; 10]); + + assert_eq!(arq.query(0, 9), 0); + + arq.update(2, 4, &-5); + arq.update(5, 7, &-3); + arq.update(1, 6, &1); + + assert_eq!(arq.query(0, 9), -3); + } + + #[test] + fn test_dynamic_rmq() { + let mut arq = DynamicArq::::new(false); + let view = arq.build_from_slice(&[0; 10]); + + assert_eq!(arq.query(view, 0, 9), 0); + + arq.update(view, 2, 4, &-5); + arq.update(view, 5, 7, &-3); + arq.update(view, 1, 6, &1); + + assert_eq!(arq.query(view, 0, 9), -3); + } + + #[test] + fn test_persistent_rmq() { + let mut arq = DynamicArq::::new(true); + let mut view = arq.build_from_slice(&[0; 10]); + + let at_init = view; + view = arq.update(view, 2, 4, &-5); + let snapshot = view; + view = arq.update(view, 5, 7, &-3); + view = arq.update(view, 1, 6, &1); + + assert_eq!(arq.query(at_init, 0, 9), 0); + assert_eq!(arq.query(snapshot, 0, 9), -5); + assert_eq!(arq.query(view, 0, 9), -3); + } + + #[test] + fn test_huge_rmq() { + let quintillion = 1_000_000_000_000_000_000; + let mut arq = DynamicArq::::new(false); + let view = arq.build_from_identity(9 * quintillion + 1); + + arq.update(view, 2 * quintillion, 4 * quintillion, &-5); + arq.update(view, 5 * quintillion, 7 * quintillion, &-3); + arq.update(view, 1 * quintillion, 6 * quintillion, &1); + + assert_eq!(arq.query(view, 0, 9 * quintillion), -3); + } + + #[test] + fn test_range_sum() { + let mut arq = StaticArq::::new(&[0; 10]); + + assert_eq!(arq.query(0, 9), 0); + + arq.update(1, 3, &10); + arq.update(3, 5, &1); + + assert_eq!(arq.query(0, 9), 23); + assert_eq!(arq.query(10, 4), 0); + } + + #[test] + fn test_dynamic_range_sum() { + let mut arq = DynamicArq::::new(false); + let view = arq.build_from_slice(&[0; 10]); + + assert_eq!(arq.query(view, 0, 9), 0); + + arq.update(view, 1, 3, &10); + arq.update(view, 3, 5, &1); + + assert_eq!(arq.query(view, 0, 9), 23); + assert_eq!(arq.query(view, 10, 4), 0); + } + + #[test] + fn test_supply_demand() { + let mut arq = StaticArq::::new(&[(0, 0, 0); 10]); + + arq.update(1, 1, &(25, 100)); + arq.update(3, 3, &(100, 30)); + arq.update(9, 9, &(0, 20)); + + assert_eq!(arq.query(0, 9), (125, 150, 75)); + } + + #[test] + fn test_dynamic_supply_demand() { + let mut arq = DynamicArq::::new(false); + let view = arq.build_from_identity(10); + + arq.update(view, 1, 1, &(25, 100)); + arq.update(view, 3, 3, &(100, 30)); + arq.update(view, 9, 9, &(0, 20)); + + assert_eq!(arq.query(view, 0, 9), (125, 150, 75)); + } + + #[test] + fn test_binary_search_rmq() { + let vec = vec![2, 1, 0, -1, -2, -3, -4, -5]; + let mut arq = StaticArq::::new(&vec); + let first_neg = static_arq::first_negative(&mut arq); + + arq.update(3, 7, &0); + let first_neg_zeros = static_arq::first_negative(&mut arq); + + assert_eq!(first_neg, Some(3)); + assert_eq!(first_neg_zeros, None); + } + + #[test] + fn test_dynamic_binary_search_rmq() { + let vec = vec![2, 1, 0, -1, -2, -3, -4, -5]; + let mut arq = DynamicArq::::new(false); + let view = arq.build_from_slice(&vec); + let first_neg = dynamic_arq::first_negative(&mut arq, view); + + arq.update(view, 3, 7, &0); + let first_neg_zeros = dynamic_arq::first_negative(&mut arq, view); + + assert_eq!(first_neg, Some(3)); + assert_eq!(first_neg_zeros, None); + } +} diff --git a/src/range_query/specs.rs b/src/range_query/specs.rs new file mode 100644 index 0000000..9a70411 --- /dev/null +++ b/src/range_query/specs.rs @@ -0,0 +1,126 @@ +//! A collection of example ArqSpec implementations + +pub trait ArqSpec { + /// Type of underlying array elements. + type S: Clone; + /// Type of data representing an endomorphism. + // Note that while a Fn(S) -> S may seem like a more natural representation + // for an endomorphism, compositions would then have to delegate to each of + // their parts. This representation is more efficient. + type F: Clone; + + /// Must satisfy the Associative Law: + /// For all a,b,c, op(a, op(b, c)) = op(op(a, b), c) + fn op(a: &Self::S, b: &Self::S) -> Self::S; + /// Must satisfy the Identity Law: + /// For all a, op(a, identity()) = op(identity(), a) = a + fn identity() -> Self::S; + /// Must satisfy the Composition Law: + /// For all f,g,a, apply(compose(f, g), a) = apply(f, apply(g, a)) + fn compose(f: &Self::F, g: &Self::F) -> Self::F; + /// Must satisfy the Distributive Law: + /// For all f,a,b, apply(f, op(a, b), s+t) = op(apply(f, a, s), apply(f, b, t)) + /// The `size` parameter makes this law easier to satisfy in certain cases. + fn apply(f: &Self::F, a: &Self::S, size: i64) -> Self::S; + + // The following relaxations to the laws may apply. + // If only point updates are made, the Composition and Distributive Laws + // no longer apply. + // - compose() is never called, so it can be left unimplemented!(). + // - apply() is only ever called on leaves, i.e., with size == 1. + // If only point queries are made, the Associative and Distributive Laws + // no longer apply. + // - op()'s result only matters when identity() is an argument. + // - apply()'s result only matters on leaves, i.e., with size == 1. +} + +/// Range Minimum Query (RMQ), a classic application of ARQ. +/// update(l, r, &f) sets all entries a[l..=r] to f. +/// query(l, r) finds the minimum value in a[l..=r]. +// +// Exercises: try augmenting this struct to find the index of a minimum element +// in a range query, as well as the number of elements equal to the minimum. +// Then instead of overwriting values with a constant assignment a[i] = f, +// try supporting addition: a[i] += f. +pub enum AssignMin {} +impl ArqSpec for AssignMin { + type S = i64; + type F = i64; + fn op(&a: &Self::S, &b: &Self::S) -> Self::S { + a.min(b) + } + fn identity() -> Self::S { + i64::MAX + } + fn compose(&f: &Self::F, _: &Self::F) -> Self::F { + f + } + fn apply(&f: &Self::F, _: &Self::S, _: i64) -> Self::S { + f + } +} + +/// Range Sum Query, a slightly trickier classic application of ARQ. +/// update(l, r, &f) sets all entries a[l..=r] to f. +/// query(l, r) sums all the entries a[l..=r]. +/// +/// # Panics +/// +/// Associated functions will panic on overflow. +// +// Note that while the `size` parameter seems necessary to satisfy the +// Distributive Law, it is merely a convenience: in essence what we've done +// is move to the product monoid of tuples (value, size_of_subtree). +// +// In mathematical jargon, we say that constant assignment f(a) = f is not an +// endomorphism on (i64, +) because f(a+b) = f != 2*f = f(a) + f(b). +// On the other hand, f((a, s)) = (f*s, s) is indeed an endomorphism on pairs +// with vector addition: f((a, s) + (b, t)) = f((a+b, s+t)) = (f*(s+t), s+t) +// = (f*s, s) + (f*t, t) = f((a,s)) + f((b,t)). +pub enum AssignSum {} +impl ArqSpec for AssignSum { + type S = i64; + type F = i64; + fn op(&a: &Self::S, &b: &Self::S) -> Self::S { + a + b + } + fn identity() -> Self::S { + 0 + } + fn compose(&f: &Self::F, _: &Self::F) -> Self::F { + f + } + fn apply(&f: &Self::F, _: &Self::S, size: i64) -> Self::S { + f * size + } +} + +/// Supply & Demand, based on https://codeforces.com/gym/102218/problem/F +/// update(i, i, &(p, o)) increases supply by p and demand by o at time i. +/// query(l, r) computes total supply and demand at times l to r, as well as +// how much of the supply is subsequently met by the demand. +// +// Note that the apply() operation is only correct when applied to leaf nodes. +// Therefore, update() must only be used in "eager" mode, i.e., with l == r. +// compose() should be unimplemented!() to prevent accidental "lazy" updates. +pub enum SupplyDemand {} +impl ArqSpec for SupplyDemand { + type S = (i64, i64, i64); // production, orders, sales + type F = (i64, i64); + fn op((p1, o1, s1): &Self::S, (p2, o2, s2): &Self::S) -> Self::S { + let extra = (p1 - s1).min(o2 - s2); + (p1 + p2, o1 + o2, s1 + s2 + extra) + } + fn identity() -> Self::S { + (0, 0, 0) + } + fn compose(_: &Self::F, _: &Self::F) -> Self::F { + unimplemented!() + } + fn apply(&(p_add, o_add): &Self::F, &(p, o, _): &Self::S, s: i64) -> Self::S { + assert_eq!(s, 1); + let p = p + p_add; + let o = o + o_add; + (p, o, p.min(o)) + } +} diff --git a/src/range_query/sqrt_decomp.rs b/src/range_query/sqrt_decomp.rs new file mode 100644 index 0000000..9accc7e --- /dev/null +++ b/src/range_query/sqrt_decomp.rs @@ -0,0 +1,117 @@ +/// A generic implementation of Mo's algorithm, aka Query Sqrt Decomposition. +/// It answers q offline queries over intervals in 0..n by shifting the query +/// interval's endpoints by one position at a time. +/// Each endpoint incurs a total cost of at most n * sqrt(q * L_OP * R_OP). +pub trait MoState { + type Q; + type A; + + /// cost ratio L_OP / R_OP between a left endpoint and a right endpoint move + const L_R_RATIO: f64 = 1.0; + + fn query(&self, q: &Self::Q) -> Self::A; + fn insert_left(&mut self, pos: usize); + fn remove_left(&mut self, pos: usize); + + fn insert_right(&mut self, pos: usize) { + self.insert_left(pos); + } + fn remove_right(&mut self, pos: usize) { + self.remove_left(pos); + } + /// After initializing self to a state corresponding to an empty interval, + /// call this function to answer all your queries. + fn process(&mut self, queries: &[(usize, usize, Self::Q)]) -> Vec { + let q = queries.len(); + let mut q_positions: Vec = (0..q).collect(); + if let Some(max_r) = queries.iter().map(|&(_, r, _)| r).max() { + let q_adjusted = q as f64 * Self::L_R_RATIO; + let bucket_width = 1 + max_r / q_adjusted.sqrt() as usize; + q_positions.sort_unstable_by_key(|&i| { + let (l, mut r) = (queries[i].0, queries[i].1); + let bucket = l / bucket_width; + if bucket % 2 != 0 { + r = max_r - r; + } + (bucket, r) + }); + } + + let (mut cur_l, mut cur_r) = (1, 0); + let mut answers = Vec::with_capacity(queries.len()); + for i in q_positions { + let (l, r, ref q) = queries[i]; + while cur_l > l { + cur_l -= 1; + self.insert_left(cur_l); + } + while cur_r < r { + cur_r += 1; + self.insert_right(cur_r); + } + while cur_l < l { + self.remove_left(cur_l); + cur_l += 1; + } + while cur_r > r { + self.remove_right(cur_r); + cur_r -= 1; + } + answers.push((i, self.query(q))); + } + answers.sort_unstable_by_key(|&(i, _)| i); + answers.into_iter().map(|(_, ans)| ans).collect() + } +} + +pub struct DistinctVals { + vals: Vec, + counts: Vec, + distinct: usize, +} +impl DistinctVals { + pub fn new(vals: Vec) -> Self { + let &max_val = vals.iter().max().unwrap_or(&0); + Self { + vals, + counts: vec![0; max_val + 1], + distinct: 0, + } + } +} +impl MoState for DistinctVals { + type Q = (); + type A = usize; + fn query(&self, _: &Self::Q) -> Self::A { + self.distinct + } + fn insert_left(&mut self, pos: usize) { + let v = self.vals[pos]; + if self.counts[v] == 0 { + self.distinct += 1; + } + self.counts[v] += 1; + } + fn remove_left(&mut self, pos: usize) { + let v = self.vals[pos]; + self.counts[v] -= 1; + if self.counts[v] == 0 { + self.distinct -= 1; + } + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_mos_algorithm() { + let queries = vec![(0, 2, ()), (5, 5, ()), (2, 6, ()), (0, 6, ())]; + let arr = vec![4, 8, 4, 7, 1, 9, 8]; + + let answers = DistinctVals::new(arr).process(&queries); + + assert_eq!(answers, vec![2, 1, 5, 5]); + } +} diff --git a/src/range_query/static_arq.rs b/src/range_query/static_arq.rs new file mode 100644 index 0000000..461986d --- /dev/null +++ b/src/range_query/static_arq.rs @@ -0,0 +1,158 @@ +//! Associative Range Query Tree +use super::ArqSpec; + +/// Colloquially known as a "segtree" in the sport programming literature, it +/// represents a sequence of elements a_i (0 <= i < size) from a monoid (S, +) +/// on which we want to support fast range operations: +/// +/// - update(l, r, f) replaces a_i (l <= i <= r) by f(a_i) for an endomorphism f +/// - query(l, r) returns the aggregate a_l + a_{l+1} + ... + a_r +/// +/// This compact representation is based on a [blog post by Al.Cash] +/// (http://codeforces.com/blog/entry/18051). All nodes have 0 or 2 children. +/// Hence, trees whose size is not a power of two will have multiple roots. +/// +/// Future work: ArqTree would lend itself naturally to Rust's ownership system. +/// Initially, we should only have access to the root nodes: +/// if size is a power of two, there is a unique root at index 1. +/// arq.push(i) locks i and acquires access to its children. +/// arq.pull(i) is called when the lock on i is released. +pub struct StaticArq { + val: Vec, + app: Vec>, +} + +impl StaticArq { + /// Initializes a static balanced binary tree on top of the given sequence. + pub fn new(init_val: &[T::S]) -> Self { + let size = init_val.len(); + let mut val = vec![T::identity(); size]; + val.extend_from_slice(init_val); + let app = vec![None; size]; + + let mut arq = Self { val, app }; + for p in (0..size).rev() { + arq.pull(p); + } + arq + } + + fn apply(&mut self, p: usize, f: &T::F, s: i64) { + self.val[p] = T::apply(f, &self.val[p], s); + if let Some(lazy) = self.app.get_mut(p) { + let h = match *lazy { + Some(ref g) => T::compose(f, g), + None => f.clone(), + }; + *lazy = Some(h); + } + } + + fn push(&mut self, p: usize) { + if let Some(ref f) = self.app[p].take() { + let s = (self.app.len().div_ceil(p) / 2).next_power_of_two() as i64; + self.apply(p << 1, f, s); + self.apply((p << 1) | 1, f, s); + } + } + + fn pull(&mut self, p: usize) { + self.val[p] = T::op(&self.val[p << 1], &self.val[(p << 1) | 1]); + } + + fn push_to(&mut self, p: usize) { + let one_plus_floor_log_p = (p + 1).next_power_of_two().trailing_zeros(); + for i in (1..one_plus_floor_log_p).rev() { + self.push(p >> i); + } + } + + fn pull_from(&mut self, mut p: usize) { + while p > 1 { + p >>= 1; + self.pull(p); + } + } + + /// Applies the endomorphism f to all entries from l to r, inclusive. + /// If l == r, the updates are eager. Otherwise, they are lazy. + /// + /// # Panics + /// + /// Panics if r >= size. Note that l > r is valid, meaning an empty range. + pub fn update(&mut self, mut l: usize, mut r: usize, f: &T::F) { + l += self.app.len(); + r += self.app.len(); + if l < r { + self.push_to(l); + } + self.push_to(r); + let (mut l0, mut r0, mut s) = (1, 1, 1); + while l <= r { + if l & 1 == 1 { + self.apply(l, f, s); + l0 = l0.max(l); + l += 1; + } + if r & 1 == 0 { + self.apply(r, f, s); + r0 = r0.max(r); + r -= 1; + } + l >>= 1; + r >>= 1; + s <<= 1; + } + self.pull_from(l0); + self.pull_from(r0); + } + + /// Returns the aggregate range query on all entries from l to r, inclusive. + /// + /// # Panics + /// + /// Panics if r >= size. Note that l > r is valid, meaning an empty range. + pub fn query(&mut self, mut l: usize, mut r: usize) -> T::S { + l += self.app.len(); + r += self.app.len(); + if l < r { + self.push_to(l); + } + self.push_to(r); + let (mut l_agg, mut r_agg) = (T::identity(), T::identity()); + while l <= r { + if l & 1 == 1 { + l_agg = T::op(&l_agg, &self.val[l]); + l += 1; + } + if r & 1 == 0 { + r_agg = T::op(&self.val[r], &r_agg); + r -= 1; + } + l >>= 1; + r >>= 1; + } + T::op(&l_agg, &r_agg) + } +} + +/// An example of binary search to find the first position whose element is negative. +/// In this case, we use RMQ to locate the leftmost negative element. +/// To ensure the existence of a valid root note (i == 1) from which to descend, +/// the tree's size must be a power of two. +pub fn first_negative(arq: &mut StaticArq) -> Option { + assert!(arq.app.len().is_power_of_two()); + let mut p = 1; + if arq.val[p] >= 0 { + None + } else { + while p < arq.app.len() { + arq.push(p); + p <<= 1; + if arq.val[p] >= 0 { + p |= 1; + } + } + Some(p - arq.app.len()) + } +} diff --git a/src/rng.rs b/src/rng.rs new file mode 100644 index 0000000..d25adf0 --- /dev/null +++ b/src/rng.rs @@ -0,0 +1,104 @@ +//! Pseudorandom number generators (PRNGs). + +/// A simple and efficient random number generator. +pub type SmallRng = Xoshiro256PlusPlus; + +/// A xoshiro256++ random number generator. +/// +/// This is a simplified version of the `SmallRng` implementation from the +/// excellent `rand` crate, keeping only essential features. +/// +/// The xoshiro256++ algorithm is not suitable for cryptographic purposes, but +/// is very fast and has excellent statistical properties. +/// +/// * Source: [Docs.rs](https://docs.rs/rand/0.8.4/src/rand/rngs/xoshiro256plusplus.rs.html) +/// * Theory: [Xorshift - Wikipedia](https://en.wikipedia.org/wiki/Xorshift) +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Xoshiro256PlusPlus { + s: [u64; 4], +} + +impl Xoshiro256PlusPlus { + /// Construct a new RNG from a 64-bit seed. + pub fn new(mut state: u64) -> Self { + const PHI: u64 = 0x9e3779b97f4a7c15; + let mut seed = <[u64; 4]>::default(); + for chunk in &mut seed { + state = state.wrapping_add(PHI); + let mut z = state; + z = (z ^ (z >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); + z = (z ^ (z >> 27)).wrapping_mul(0x94d049bb133111eb); + z = z ^ (z >> 31); + *chunk = z; + } + Self { s: seed } + } + + /// Generate a random `u32`. + #[inline] + pub fn next_u32(&mut self) -> u32 { + (self.next_u64() >> 32) as u32 + } + + /// Generate a random `u64`. + #[inline] + pub fn next_u64(&mut self) -> u64 { + let result_plusplus = self.s[0] + .wrapping_add(self.s[3]) + .rotate_left(23) + .wrapping_add(self.s[0]); + + let t = self.s[1] << 17; + + self.s[2] ^= self.s[0]; + self.s[3] ^= self.s[1]; + self.s[1] ^= self.s[2]; + self.s[0] ^= self.s[3]; + + self.s[2] ^= t; + + self.s[3] = self.s[3].rotate_left(45); + + result_plusplus + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_xoshiro256plusplus() { + let mut rng = Xoshiro256PlusPlus::new(42); + assert_eq!(rng.next_u64(), 15021278609987233951); + assert_eq!(rng.next_u64(), 5881210131331364753); + assert_eq!(rng.next_u64(), 18149643915985481100); + assert_eq!(rng.next_u64(), 12933668939759105464); + assert_eq!(rng.next_u64(), 14637574242682825331); + assert_eq!(rng.next_u64(), 10848501901068131965); + assert_eq!(rng.next_u64(), 2312344417745909078); + assert_eq!(rng.next_u64(), 11162538943635311430); + } + + #[test] + fn reference() { + let mut rng = Xoshiro256PlusPlus { s: [1, 2, 3, 4] }; + // These values were produced with the reference implementation: + // http://xoshiro.di.unimi.it/xoshiro256plusplus.c + let expected = [ + 41943041, + 58720359, + 3588806011781223, + 3591011842654386, + 9228616714210784205, + 9973669472204895162, + 14011001112246962877, + 12406186145184390807, + 15849039046786891736, + 10450023813501588000, + ]; + for &e in &expected { + assert_eq!(rng.next_u64(), e); + } + } +} diff --git a/src/scanner.rs b/src/scanner.rs index cb49dfc..83fb816 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -1,65 +1,136 @@ //! Generic utility for reading data from standard input, based on [voxl's //! stdin wrapper](http://codeforces.com/contest/702/submission/19589375). use std::io; +use std::str; /// Reads white-space separated tokens one at a time. -pub struct Scanner { - reader: B, +pub struct Scanner { + reader: R, buffer: Vec, } -impl Scanner { - pub fn new(reader: B) -> Self { +impl Scanner { + pub fn new(reader: R) -> Self { Self { - reader: reader, - buffer: Vec::new(), + reader, + buffer: vec![], } } - /// Use "turbofish" syntax next::() to select data type of next token. - pub fn next(&mut self) -> T - where - T::Err: ::std::fmt::Debug, - { - if let Some(front) = self.buffer.pop() { - front.parse::().expect(&front) - } else { + /// Use "turbofish" syntax token::() to select data type of next token. + /// + /// # Panics + /// + /// Panics if there's an I/O error or if the token cannot be parsed as T. + pub fn token(&mut self) -> T { + loop { + if let Some(token) = self.buffer.pop() { + return token.parse().ok().expect("Failed parse"); + } let mut input = String::new(); - self.reader.read_line(&mut input).expect("Line not read"); + self.reader.read_line(&mut input).expect("Failed read"); self.buffer = input.split_whitespace().rev().map(String::from).collect(); - self.next() } } } +/// Same API as Scanner but nearly twice as fast, using horribly unsafe dark arts +pub struct UnsafeScanner { + reader: R, + buf_str: Vec, + buf_iter: str::SplitAsciiWhitespace<'static>, +} + +impl UnsafeScanner { + pub fn new(reader: R) -> Self { + Self { + reader, + buf_str: vec![], + buf_iter: "".split_ascii_whitespace(), + } + } + + /// This function should be marked unsafe, but noone has time for that in a + /// programming contest. Use at your own risk! + pub fn token(&mut self) -> T { + loop { + if let Some(token) = self.buf_iter.next() { + return token.parse().ok().expect("Failed parse"); + } + self.buf_str.clear(); + self.reader + .read_until(b'\n', &mut self.buf_str) + .expect("Failed read"); + self.buf_iter = unsafe { + let slice = str::from_utf8_unchecked(&self.buf_str); + std::mem::transmute(slice.split_ascii_whitespace()) + } + } + } +} + +pub fn scanner_from_file(filename: &str) -> Scanner> { + let file = std::fs::File::open(filename).expect("Input file not found"); + Scanner::new(io::BufReader::new(file)) +} + +pub fn writer_to_file(filename: &str) -> io::BufWriter { + let file = std::fs::File::create(filename).expect("Output file not found"); + io::BufWriter::new(file) +} #[cfg(test)] mod test { use super::*; + fn solve(scan: &mut Scanner, out: &mut W) { + let x = scan.token::(); + let y = scan.token::(); + writeln!(out, "{} - {} = {}", x, y, x - y).ok(); + } + + fn unsafe_solve(scan: &mut UnsafeScanner, out: &mut W) { + let x = scan.token::(); + let y = scan.token::(); + writeln!(out, "{} - {} = {}", x, y, x - y).ok(); + } + + #[test] + fn test_in_memory_io() { + let input: &[u8] = b"50 8"; + let mut scan = Scanner::new(input); + let mut out = vec![]; + + solve(&mut scan, &mut out); + assert_eq!(out, b"50 - 8 = 42\n"); + } + #[test] - fn test_fake_input() { - let cursor = io::Cursor::new("44 2"); - let mut scan = Scanner::new(cursor); - let x = scan.next::(); - let y = scan.next::(); - assert_eq!(x - y, 42); + fn test_in_memory_unsafe() { + let input: &[u8] = b"50 8"; + let mut scan = UnsafeScanner::new(input); + let mut out = vec![]; + + unsafe_solve(&mut scan, &mut out); + assert_eq!(out, b"50 - 8 = 42\n"); } #[test] - fn test_stdin() { - let stdin = io::stdin(); - let mut scan = Scanner::new(stdin.lock()); + fn test_compile_stdio() { + let mut scan = Scanner::new(io::stdin().lock()); + let mut out = io::BufWriter::new(io::stdout().lock()); + if false { - let _ = scan.next::(); + solve(&mut scan, &mut out); } } #[test] - #[should_panic] - fn test_file() { - let file = ::std::fs::File::open("asdf.txt").expect("File not found"); - let mut scan = Scanner::new(io::BufReader::new(file)); - let _ = scan.next::(); + #[should_panic(expected = "Input file not found")] + fn test_panic_file() { + let mut scan = scanner_from_file("input_file.txt"); + let mut out = writer_to_file("output_file.txt"); + + solve(&mut scan, &mut out); } } diff --git a/src/string_proc.rs b/src/string_proc.rs index bd026ab..4716f68 100644 --- a/src/string_proc.rs +++ b/src/string_proc.rs @@ -1,81 +1,359 @@ //! String processing algorithms. +use std::cmp::{max, min}; +use std::collections::{HashMap, VecDeque, hash_map::Entry}; -/// Data structure for Knuth-Morris-Pratt string matching against a pattern. -pub struct Matcher<'a> { - pub pattern: &'a [u8], +/// Prefix trie, easily augmentable by adding more fields and/or methods +pub struct Trie { + links: Vec>, +} + +impl Default for Trie { + /// Creates an empty trie with a root node. + fn default() -> Self { + Self { + links: vec![HashMap::new()], + } + } +} + +impl Trie { + /// Inserts a word into the trie, and returns the index of its node. + pub fn insert(&mut self, word: impl IntoIterator) -> usize { + let mut node = 0; + + for ch in word { + let len = self.links.len(); + node = match self.links[node].entry(ch) { + Entry::Occupied(entry) => *entry.get(), + Entry::Vacant(entry) => { + entry.insert(len); + self.links.push(HashMap::new()); + len + } + } + } + node + } + + /// Finds a word in the trie, and returns the index of its node. + pub fn get(&self, word: impl IntoIterator) -> Option { + let mut node = 0; + for ch in word { + node = *self.links[node].get(&ch)?; + } + Some(node) + } +} + +/// Single-pattern matching with the Knuth-Morris-Pratt algorithm +pub struct Matcher<'a, C: Eq> { + /// The string pattern to search for. + pub pattern: &'a [C], + /// KMP match failure automaton: fail[i] is the length of the longest + /// string that's both a proper prefix and a proper suffix of pattern[0..=i]. pub fail: Vec, } -impl<'a> Matcher<'a> { - /// Sets fail[i] = length of longest proper prefix-suffix of pattern[0...i]. - pub fn new(pattern: &'a [u8]) -> Self { +impl<'a, C: Eq> Matcher<'a, C> { + /// Precomputes the automaton that allows linear-time string matching. + /// + /// # Example + /// + /// ``` + /// use contest_algorithms::string_proc::Matcher; + /// let byte_string: &[u8] = b"hello"; + /// let utf8_string: &str = "hello"; + /// let vec_char: Vec = utf8_string.chars().collect(); + /// + /// let match_from_byte_literal = Matcher::new(byte_string); + /// let match_from_utf8 = Matcher::new(utf8_string.as_bytes()); + /// let match_from_chars = Matcher::new(&vec_char); + /// + /// let vec_int = vec![4, -3, 1]; + /// let match_from_ints = Matcher::new(&vec_int); + /// ``` + /// + /// # Panics + /// + /// Panics if pattern is empty. + pub fn new(pattern: &'a [C]) -> Self { let mut fail = Vec::with_capacity(pattern.len()); fail.push(0); let mut len = 0; - for &ch in &pattern[1..] { - while len > 0 && pattern[len] != ch { + for ch in &pattern[1..] { + while len > 0 && pattern[len] != *ch { len = fail[len - 1]; } - if pattern[len] == ch { + if pattern[len] == *ch { len += 1; } fail.push(len); } + Self { pattern, fail } + } + + /// KMP algorithm, sets @return[i] = length of longest prefix of pattern + /// matching a suffix of text[0..=i]. + pub fn kmp_match(&self, text: impl IntoIterator) -> Vec { + let mut len = 0; + text.into_iter() + .map(|ch| { + if len == self.pattern.len() { + len = self.fail[len - 1]; + } + while len > 0 && self.pattern[len] != ch { + len = self.fail[len - 1]; + } + if self.pattern[len] == ch { + len += 1; + } + len + }) + .collect() + } +} + +/// Multi-pattern matching with the Aho-Corasick algorithm +pub struct MultiMatcher { + /// A prefix trie storing the string patterns to search for. + pub trie: Trie, + /// Stores which completed pattern string each node corresponds to. + pub pat_id: Vec>, + /// Aho-Corasick failure automaton. fail[i] is the node corresponding to the + /// longest prefix-suffix of the node corresponding to i. + pub fail: Vec, + /// Shortcut to the next match along the failure chain, or to the root. + pub fast: Vec, +} + +impl MultiMatcher { + fn next(trie: &Trie, fail: &[usize], mut node: usize, ch: &C) -> usize { + loop { + if let Some(&child) = trie.links[node].get(ch) { + return child; + } else if node == 0 { + return 0; + } + node = fail[node]; + } + } + + /// Precomputes the automaton that allows linear-time string matching. + /// If there are duplicate patterns, all but one copy will be ignored. + pub fn new(patterns: impl IntoIterator>) -> Self { + let mut trie = Trie::default(); + #[allow(clippy::needless_collect)] // It's not needless: it affects trie.links.len() + let pat_nodes: Vec = patterns.into_iter().map(|pat| trie.insert(pat)).collect(); + + let mut pat_id = vec![None; trie.links.len()]; + for (i, node) in pat_nodes.into_iter().enumerate() { + pat_id[node] = Some(i); + } + + let mut fail = vec![0; trie.links.len()]; + let mut fast = vec![0; trie.links.len()]; + let mut q: VecDeque = trie.links[0].values().cloned().collect(); + + while let Some(node) = q.pop_front() { + for (ch, &child) in &trie.links[node] { + let nx = Self::next(&trie, &fail, fail[node], ch); + fail[child] = nx; + fast[child] = if pat_id[nx].is_some() { nx } else { fast[nx] }; + q.push_back(child); + } + } + Self { - pattern: pattern, - fail: fail, + trie, + pat_id, + fail, + fast, } } - /// KMP algorithm, sets matches[i] = length of longest prefix of pattern - /// matching a suffix of text[0...i]. - pub fn kmp_match(&self, text: &[u8]) -> Vec { - let mut matches = Vec::with_capacity(text.len()); - let mut len = 0; - for &ch in text { - if len == self.pattern.len() { - len = self.fail[len - 1]; + /// Aho-Corasick algorithm, sets @return[i] = node corresponding to + /// longest prefix of some pattern matching a suffix of text[0..=i]. + pub fn ac_match(&self, text: impl IntoIterator) -> Vec { + let mut node = 0; + text.into_iter() + .map(|ch| { + node = Self::next(&self.trie, &self.fail, node, &ch); + node + }) + .collect() + } + + /// For each non-empty match, returns where in the text it ends, and the index + /// of the corresponding pattern. + pub fn get_end_pos_and_pat_id(&self, match_nodes: &[usize]) -> Vec<(usize, usize)> { + let mut res = vec![]; + for (text_pos, &(mut node)) in match_nodes.iter().enumerate() { + while node != 0 { + if let Some(id) = self.pat_id[node] { + res.push((text_pos + 1, id)); + } + node = self.fast[node]; } - while len > 0 && self.pattern[len] != ch { - len = self.fail[len - 1]; + } + res + } +} + +/// Suffix array data structure, useful for a variety of string queries. +pub struct SuffixArray { + /// The suffix array itself, holding suffix indices in sorted order. + pub sfx: Vec, + /// rank[i][j] = rank of the j'th suffix, considering only 2^i chars. + /// In other words, rank[i] is a ranking of the substrings text[j..j+2^i]. + pub rank: Vec>, +} + +impl SuffixArray { + /// O(n + max_key) stable sort on the items generated by vals. + /// Items v in vals are sorted according to val_to_key[v]. + fn counting_sort( + vals: impl Iterator + Clone, + val_to_key: &[usize], + max_key: usize, + ) -> Vec { + let mut counts = vec![0; max_key]; + for v in vals.clone() { + counts[val_to_key[v]] += 1; + } + let mut total = 0; + for c in counts.iter_mut() { + total += *c; + *c = total - *c; + } + let mut result = vec![0; total]; + for v in vals { + let c = &mut counts[val_to_key[v]]; + result[*c] = v; + *c += 1; + } + result + } + + /// Suffix array construction in O(n log n) time. + pub fn new(text: impl IntoIterator) -> Self { + let init_rank = text.into_iter().map(|ch| ch as usize).collect::>(); + let n = init_rank.len(); + let mut sfx = Self::counting_sort(0..n, &init_rank, 256); + let mut rank = vec![init_rank]; + // Invariant at the start of every loop iteration: + // suffixes are sorted according to the first skip characters. + for skip in (0..).map(|i| 1 << i).take_while(|&skip| skip < n) { + let prev_rank = rank.last().unwrap(); + let mut cur_rank = prev_rank.clone(); + + let pos = (n - skip..n).chain(sfx.into_iter().filter_map(|p| p.checked_sub(skip))); + sfx = Self::counting_sort(pos, prev_rank, max(n, 256)); + + let mut prev = sfx[0]; + cur_rank[prev] = 0; + for &cur in sfx.iter().skip(1) { + if max(prev, cur) + skip < n + && prev_rank[prev] == prev_rank[cur] + && prev_rank[prev + skip] == prev_rank[cur + skip] + { + cur_rank[cur] = cur_rank[prev]; + } else { + cur_rank[cur] = cur_rank[prev] + 1; + } + prev = cur; } - if self.pattern[len] == ch { - len += 1; + rank.push(cur_rank); + } + Self { sfx, rank } + } + + /// Computes the length of longest common prefix of text[i..] and text[j..]. + pub fn longest_common_prefix(&self, mut i: usize, mut j: usize) -> usize { + let mut len = 0; + for (k, rank) in self.rank.iter().enumerate().rev() { + if rank[i] == rank[j] { + i += 1 << k; + j += 1 << k; + len += 1 << k; + if max(i, j) >= self.sfx.len() { + break; + } } - matches.push(len); } - matches + len } } /// Manacher's algorithm for computing palindrome substrings in linear time. -/// len[2*i] = odd length of palindrome centred at text[i]. -/// len[2*i+1] = even length of palindrome centred at text[i+0.5]. -pub fn palindromes(text: &[u8]) -> Vec { - let mut len = Vec::with_capacity(2 * text.len() - 1); - len.push(1); - while len.len() < len.capacity() { - let i = len.len() - 1; - let max_len = ::std::cmp::min(i + 1, len.capacity() - i); - while len[i] < max_len && text[(i - len[i] - 1) / 2] == text[(i + len[i] + 1) / 2] { - len[i] += 2; - } - if len[i] < 2 { - let a = 1 - len[i]; - len.push(a); +/// pal[2*i] = odd length of palindrome centred at text[i]. +/// pal[2*i+1] = even length of palindrome centred at text[i+0.5]. +/// +/// # Panics +/// +/// Panics if text is empty. +pub fn palindromes(text: &[impl Eq]) -> Vec { + let mut pal = Vec::with_capacity(2 * text.len() - 1); + pal.push(1); + while pal.len() < pal.capacity() { + let i = pal.len() - 1; + let max_len = min(i + 1, pal.capacity() - i); + while pal[i] < max_len && text[(i - pal[i] - 1) / 2] == text[(i + pal[i] + 1) / 2] { + pal[i] += 2; + } + if let Some(a) = 1usize.checked_sub(pal[i]) { + pal.push(a); } else { for d in 1.. { - let (a, b) = (len[i - d], len[i] - d); + let (a, b) = (pal[i - d], pal[i] - d); if a < b { - len.push(a); + pal.push(a); } else { - len.push(b); + pal.push(b); break; } } } } - len + pal +} + +/// Z algorithm: computes the array Z[..], where Z[i] is the length of the +/// longest text prefix of text[i..] that is **also a prefix** of text. +/// +/// It runs in O(n) time, maintaining the invariant that l <= i and +/// text[0..r-l] == text[l..r]. It can be embedded in a larger algorithm, +/// or used for string searching as an alternative to KMP. +/// +/// # Example +/// +/// ``` +/// use contest_algorithms::string_proc::z_algorithm; +/// let z = z_algorithm(b"ababbababbabababbabababbababbaba"); +/// assert_eq!( +/// z, +/// vec![ +/// 32, 0, 2, 0, 0, 9, 0, 2, 0, 0, 4, 0, 9, 0, 2, 0, 0, 4, 0, 13, 0, 2, +/// 0, 0, 8, 0, 2, 0, 0, 3, 0, 1, +/// ], +/// ); +/// ``` +pub fn z_algorithm(text: &[impl Eq]) -> Vec { + let n = text.len(); + let (mut l, mut r) = (1, 1); + let mut z = Vec::with_capacity(n); + z.push(n); + for i in 1..n { + if r > i + z[i - l] { + z.push(z[i - l]); + } else { + l = i; + while r < i || (r < n && text[r - i] == text[r]) { + r += 1; + } + z.push(r - i); + } + } + z } #[cfg(test)] @@ -83,14 +361,75 @@ mod test { use super::*; #[test] - fn test_string() { - let text = "abcbc".as_bytes(); - let pattern = "bc".as_bytes(); + fn test_trie() { + let dict = vec!["banana", "benefit", "banapple", "ban"]; + + let trie = dict.into_iter().fold(Trie::default(), |mut trie, word| { + trie.insert(word.bytes()); + trie + }); + + assert_eq!(trie.get("".bytes()), Some(0)); + assert_eq!(trie.get("b".bytes()), Some(1)); + assert_eq!(trie.get("banana".bytes()), Some(6)); + assert_eq!(trie.get("be".bytes()), Some(7)); + assert_eq!(trie.get("bane".bytes()), None); + } + + #[test] + fn test_kmp_matching() { + let pattern = "ana"; + let text = "banana"; + + let matches = Matcher::new(pattern.as_bytes()).kmp_match(text.bytes()); + + assert_eq!(matches, vec![0, 1, 2, 3, 2, 3]); + } + + #[test] + fn test_ac_matching() { + let dict = vec!["banana", "benefit", "banapple", "ban", "fit"]; + let text = "banana bans, apple benefits."; + + let matcher = MultiMatcher::new(dict.iter().map(|s| s.bytes())); + let match_nodes = matcher.ac_match(text.bytes()); + let end_pos_and_id = matcher.get_end_pos_and_pat_id(&match_nodes); + + assert_eq!( + end_pos_and_id, + vec![(3, 3), (6, 0), (10, 3), (26, 1), (26, 4)] + ); + } + + #[test] + fn test_suffix_array() { + let text1 = "bobocel"; + let text2 = "banana"; + + let sfx1 = SuffixArray::new(text1.bytes()); + let sfx2 = SuffixArray::new(text2.bytes()); + + assert_eq!(sfx1.sfx, vec![0, 2, 4, 5, 6, 1, 3]); + assert_eq!(sfx2.sfx, vec![5, 3, 1, 0, 4, 2]); + + assert_eq!(sfx1.longest_common_prefix(0, 2), 2); + assert_eq!(sfx2.longest_common_prefix(1, 3), 3); + + // Check that sfx and rank.last() are essentially inverses of each other. + for (p, &r) in sfx1.rank.last().unwrap().iter().enumerate() { + assert_eq!(sfx1.sfx[r], p); + } + for (p, &r) in sfx2.rank.last().unwrap().iter().enumerate() { + assert_eq!(sfx2.sfx[r], p); + } + } + + #[test] + fn test_palindrome() { + let text = "banana"; - let matches = Matcher::new(pattern).kmp_match(text); - assert_eq!(matches, vec![0, 1, 2, 1, 2]); + let pal_len = palindromes(text.as_bytes()); - let pal_len = palindromes(text); - assert_eq!(pal_len, vec![1, 0, 1, 0, 3, 0, 3, 0, 1]); + assert_eq!(pal_len, vec![1, 0, 1, 0, 3, 0, 5, 0, 3, 0, 1]); } } diff --git a/tests/codeforces343d.rs b/tests/codeforces343d.rs new file mode 100644 index 0000000..ec0f447 --- /dev/null +++ b/tests/codeforces343d.rs @@ -0,0 +1,105 @@ +//! Solves [Water Tree](http://codeforces.com/contest/343/problem/D). +//! To make a self-contained file for contest submission, dump each desired +//! module's contents directly here instead of the use statements. +//! Also, use the commented code in main() to employ standard I/O. +extern crate contest_algorithms; +use contest_algorithms::graph::Graph; +use contest_algorithms::range_query::{StaticArq, specs::AssignSum}; +use contest_algorithms::scanner::Scanner; +use std::io; + +const SAMPLE_INPUT: &[u8] = b"\ +5 +1 2 +5 1 +2 3 +4 2 +12 +1 1 +2 3 +3 1 +3 2 +3 3 +3 4 +1 2 +2 4 +3 1 +3 3 +3 4 +3 5 +"; +const SAMPLE_OUTPUT: &[u8] = b"\ +0 +0 +0 +1 +0 +1 +0 +1 +"; + +fn dfs( + graph: &Graph, + u: usize, + l: &mut [usize], + r: &mut [usize], + p: &mut [usize], + time: &mut usize, +) { + *time += 1; + l[u] = *time; + + for (_, v) in graph.adj_list(u) { + if l[v] == 0 { + p[v] = l[u]; + dfs(graph, v, l, r, p, time); + } + } + + r[u] = *time; +} + +fn solve(scan: &mut Scanner, out: &mut W) { + let n = scan.token::(); + let mut tree = Graph::new(n, 2 * (n - 1)); + for _ in 1..n { + let u = scan.token::() - 1; + let v = scan.token::() - 1; + tree.add_undirected_edge(u, v); + } + + let mut l = vec![0; n]; + let mut r = vec![0; n]; + let mut p = vec![0; n]; + dfs(&tree, 0, &mut l, &mut r, &mut p, &mut 0); + + let mut arq = StaticArq::::new(&vec![0; n + 1]); + let q = scan.token::(); + for _ in 0..q { + let c = scan.token::(); + let v = scan.token::() - 1; + let len = (r[v] - l[v] + 1) as i64; + let sum = arq.query(l[v], r[v]); + if c == 1 { + if sum != len { + arq.update(p[v], p[v], &0); + arq.update(l[v], r[v], &1); + } + } else if c == 2 { + arq.update(l[v], l[v], &0); + } else { + let ans = if sum == len { 1 } else { 0 }; + writeln!(out, "{}", ans).ok(); + } + } +} + +#[test] +fn main() { + let mut scan = Scanner::new(SAMPLE_INPUT); + let mut out = vec![]; + solve(&mut scan, &mut out); + + assert_eq!(out, SAMPLE_OUTPUT); +} diff --git a/tests/example_problem.rs b/tests/example_problem.rs deleted file mode 100644 index 0760214..0000000 --- a/tests/example_problem.rs +++ /dev/null @@ -1,28 +0,0 @@ -// To make a self-contained file for contest submission, dump each desired -// module's contents directly here instead of these use statements. -extern crate algorithms; -use algorithms::scanner::*; -use algorithms::graph::*; - -fn main1() { - let stdin = std::io::stdin(); - let mut scan = Scanner::new(stdin.lock()); - let mut graph = Graph::new(1, 0); - for _ in 0..0 { - let u = scan.next::(); - let v = scan.next::(); - graph.add_edge(u, v); - } -} - -#[test] -fn main() { - // If your contest solution requires a lot of stack space, make sure to - // run it in a custom thread. - std::thread::Builder::new() - .stack_size(50_000_000) - .spawn(main1) - .unwrap() - .join() - .unwrap(); -}