diff --git a/solution/3200-3299/3234.Count the Number of Substrings With Dominant Ones/README.md b/solution/3200-3299/3234.Count the Number of Substrings With Dominant Ones/README.md index 2f5af3e1fb0f5..1e7d61c5d1151 100644 --- a/solution/3200-3299/3234.Count the Number of Substrings With Dominant Ones/README.md +++ b/solution/3200-3299/3234.Count the Number of Substrings With Dominant Ones/README.md @@ -333,6 +333,40 @@ function numberOfSubstrings(s: string): number { } ``` +#### Rust + +```rust +impl Solution { + pub fn number_of_substrings(s: String) -> i32 { + let n = s.len(); + let mut nxt = vec![n; n + 1]; + + for i in (0..n).rev() { + nxt[i] = nxt[i + 1]; + if &s[i..i + 1] == "0" { + nxt[i] = i; + } + } + + let mut ans = 0; + for i in 0..n { + let mut cnt0 = if &s[i..i + 1] == "0" { 1 } else { 0 }; + let mut j = i; + while j < n && (cnt0 * cnt0) as i64 <= n as i64 { + let cnt1 = nxt[j + 1] - i - cnt0; + if cnt1 >= (cnt0 * cnt0) { + ans += std::cmp::min(nxt[j + 1] - j, cnt1 - cnt0 * cnt0 + 1); + } + j = nxt[j + 1]; + cnt0 += 1; + } + } + + ans as i32 + } +} +``` + diff --git a/solution/3200-3299/3234.Count the Number of Substrings With Dominant Ones/README_EN.md b/solution/3200-3299/3234.Count the Number of Substrings With Dominant Ones/README_EN.md index f2fc538c53418..00f9cd3145fc2 100644 --- a/solution/3200-3299/3234.Count the Number of Substrings With Dominant Ones/README_EN.md +++ b/solution/3200-3299/3234.Count the Number of Substrings With Dominant Ones/README_EN.md @@ -331,6 +331,40 @@ function numberOfSubstrings(s: string): number { } ``` +#### Rust + +```rust +impl Solution { + pub fn number_of_substrings(s: String) -> i32 { + let n = s.len(); + let mut nxt = vec![n; n + 1]; + + for i in (0..n).rev() { + nxt[i] = nxt[i + 1]; + if &s[i..i + 1] == "0" { + nxt[i] = i; + } + } + + let mut ans = 0; + for i in 0..n { + let mut cnt0 = if &s[i..i + 1] == "0" { 1 } else { 0 }; + let mut j = i; + while j < n && (cnt0 * cnt0) as i64 <= n as i64 { + let cnt1 = nxt[j + 1] - i - cnt0; + if cnt1 >= (cnt0 * cnt0) { + ans += std::cmp::min(nxt[j + 1] - j, cnt1 - cnt0 * cnt0 + 1); + } + j = nxt[j + 1]; + cnt0 += 1; + } + } + + ans as i32 + } +} +``` + diff --git a/solution/3200-3299/3234.Count the Number of Substrings With Dominant Ones/Solution.rs b/solution/3200-3299/3234.Count the Number of Substrings With Dominant Ones/Solution.rs new file mode 100644 index 0000000000000..55387a9144971 --- /dev/null +++ b/solution/3200-3299/3234.Count the Number of Substrings With Dominant Ones/Solution.rs @@ -0,0 +1,29 @@ +impl Solution { + pub fn number_of_substrings(s: String) -> i32 { + let n = s.len(); + let mut nxt = vec![n; n + 1]; + + for i in (0..n).rev() { + nxt[i] = nxt[i + 1]; + if &s[i..i + 1] == "0" { + nxt[i] = i; + } + } + + let mut ans = 0; + for i in 0..n { + let mut cnt0 = if &s[i..i + 1] == "0" { 1 } else { 0 }; + let mut j = i; + while j < n && (cnt0 * cnt0) as i64 <= n as i64 { + let cnt1 = nxt[j + 1] - i - cnt0; + if cnt1 >= (cnt0 * cnt0) { + ans += std::cmp::min(nxt[j + 1] - j, cnt1 - cnt0 * cnt0 + 1); + } + j = nxt[j + 1]; + cnt0 += 1; + } + } + + ans as i32 + } +}