H - Yet Another Sigma Problem 解説 by ngtkana


\(S\) はソート済みであるとします。集合 \(T = \{{ (i, j) : 1 \le i \le N, 1 \le j \le |S_i| \}}\) 上の同値関係 \(~\)

\[ (i, j) ~ (i + 1, j) \Leftarrow S_{i,1} S_{i,2} \dots S_{i,j} = S_{i + 1,1} S_{i + 1,2} \dots S_{i + 1,j} \]

を満たす最小のものであるとすると、これにより \(T\) は次の図のように分割されます。このとき答えは

\[ \sum _ {C \in T / ~} \frac { |C| (|C| - 1) } {2} \]

になります。

そこで \(S\) を添字の昇順に見ていって、同値類の終端を見つけるたびに答えに足していきましょう。そのためには同値類の始端を見つけるたびに、文字種と開始インデックスのペアをスタックに入れておけばよいです。

提出(21 ms):https://atcoder.jp/contests/abc353/submissions/53393003

use proconio::input;

fn main() {
    input! {
        n: usize,
        mut s: [String; n],
    }
    s.sort_unstable();
    s.push(String::new());
    let mut stack = Vec::new();
    let mut ans = 0;
    for (i, s) in s.iter().enumerate() {
        let lcp = stack
            .iter()
            .zip(s.chars())
            .take_while(|&(&(c, _), d)| c == d)
            .count();
        while stack.len() > lcp {
            let (_, start) = stack.pop().unwrap();
            ans += (i - start) * (i - start - 1) / 2;
        }
        for c in s[lcp..].chars() {
            stack.push((c, i));
        }
    }
    println!("{}", ans);
}

投稿日時:
最終更新: