概要
2026/04/26のABC455 E問題について、解説ACをしました。
難しい部分がいくつかあったので、解説を作成しました。
参考
- 問題URL
- 提出URL
- 注意点
- 提出コードについて
- 利用言語はGo言語です
- Atnuhs/go-aclをライブラリとして利用し、Atnuhs/go-bundlerにて1ファイルにまとめてから提出しているため、提出コードは若干見づらいです
- 解説ACをとるためにAIを利用しています
- 最終的な解法はeditorialと同様ですが、editorialとはちがうアプローチで包除原理にたどり着きました
- 提出コードについて
解説
まず、求めたい値を考えます。$A(l,r)$をS[l:r]に含まれる"A"の個数とし、$B(l,r)$、$C(l,r)$も同様にします。すると求めたい値は、
$$
\sum_{l=0}^{N-1}\sum_{r=l+1}^{N}
\left[A(l,r) \neq B(l,r) \text{かつ} A(l,r) \neq C(l,r) \text{かつ} B(l,r) \neq C(l,r)\right]
\tag{1}
$$
と表現できます。
ここでの「$[$」と「$]」$はIverson bracketというもので、例えば$\left[A(l,r) \neq B(l,r)\right]$は以下と同値です。
\begin{cases}
1 & \text{if } A(l,r) \neq B(l,r) \\
0 & \text{otherwise}
\end{cases}
つまり、$\left[A(l,r) \neq B(l,r) \text{かつ} A(l,r) \neq C(l,r) \text{かつ} B(l,r) \neq C(l,r)\right]$は、区間$[l,r)$に含まれる"A"、"B"、"C"の出現回数が相異なる場合、1、そうでなければ0となることを意味します。最初の式ではそれをすべての$l$と$r$に対して和をとっているため、すべての$S$の部分文字列に対して上記の総数を満たす部分文字列の総数（求めたい値）を表しています。
また、$\text{かつ}$を置き換えて以下の式に変形します。
$$
\sum_{l=0}^{N-1}\sum_{r=l+1}^{N}
[A(l,r) \neq B(l,r)]\times[A(l,r) \neq C(l,r)]\times[B(l,r) \neq C(l,r)]
\tag{2}
$$
素朴に各区間 [l,r) について $A(l,r),B(l,r),C(l,r)$ を数えると、区間数が $O(N^2)$、各区間の集計が $O(N)$ なので全体で $O(N^3)$ になります。
まずは累積和を使って、各区間の文字数を O(1) で求められるようにします。
$A(l,r) = A(0,r) - A(0,l)$とすると$A(0,i)$は累積和であるため、$A(l,r)$は$O(1)$で求められるようになります。
$A(0, i)$を$C_A(i)$として、$A(l,r) \neq B(l,r)$を累積和の式に書き換えて整理すると以下が得られます。
$$
C_A(l) - C_B(l) \neq C_A(r) - C_B(r)
$$
ここで、$X(i) = C_A(i) - C_B(i)$と置くことで
$$
X(l) \neq X(r)
$$
が得られます。$Y$と$Z$についても同様に、
\begin{align}
Y(i) = C_A(i) - C_C(i) \\
Z(i) = C_B(i)-C_C(i) \\
\end{align}
と置くことで、元の式(2)から式(3)を得ることができました。
$$
\sum_{l=0}^{N-1}\sum_{r=l+1}^{N}
[X(l) \neq X(r)][Y(l) \neq Y(r)][Z(l) \neq Z(r)]
\tag{3}
$$
ここでいったん$X$のみバージョンを考えてみます。
$$
\sum_{l=0}^{N-1}\sum_{r=l+1}^{N}
\left[X(l) \neq X(r) \right]
\tag{4}
$$
こちらはそのままの計算で$O(N^2)$となります。
val := 0
for i := range n {
for j := i+1; j <= n; j++ {
if x[i] != x[j] { val++ }
}
}
余事象を用いることにより$O(N)$で求めることができます。
$[X(l) \neq X(r)]$の余事象は$1-[X(l) = X(r)]$なので余事象を使って変形すると式(4)は次のようになります。
\begin{align}
&\sum_{l=0}^{N-1}\sum_{r=l+1}^{N}
(1-[X(l) = X(r)]) \\
&=\sum_{l=0}^{N-1}\sum_{r=l+1}^{N}1 - \sum_{l=0}^{N-1}\sum_{r=l+1}^{N}[X(l) = X(r)]\\
&=\sum_{l<r} 1 - \sum_{l<r}[X(l) = X(r)]\\
&=\frac{n(n+1)}{2} - \sum_{l<r}[X(l) = X(r)]
\end{align}
$\sum_{l=0}^{N-1}\sum_{r=l+1}^{N}$を省略して$\sum_{l<r}$としています。
連想配列と組み合わせを用いて、以下のように求めることができます(sortを利用してもよいです)。
// 連想配列 cnt を用いて、各値 v について x_i = v となる i の個数を数える。
cnt := make(map[int]int)
for i := range n+1 {
cnt[x[i]]++
}
// v個出現している値から二つを選ぶ選び方の個数
ng := 0
for _, v := range cnt {
ng += (v*(v-1))/2
}
total := (n*(n+1))/2
ans := total - ng
以上により$X$のみバージョンの式(4)について、$O(N)$で求めることができました。
元の式(3)に戻ります。
$[X(l) \neq X(r)][Y(l) \neq Y(r)][Z(l) \neq Z(r)]$についても同様に余事象を適用することで計算量を改善することができるのでしょうか。
式(3)余事象で置き換えた式を展開して出てくる各項を想像してみると、$[X(l)= X(r)][Y(l) = Y(r)]$や、$[X(l)= X(r)][Y(l) = Y(r)][Z(l) = Z(r)]$という項が出てきます。
これらは、連想配列のkeyを各$X$,$Y$,$Z$の値のタプル値とすることで$[X(l)=X(r)$と同様に計算量を改善することができます。
// goの場合、連想配列のkeyとして固定配列が利用可能
cnt := make(map[[3]int]int)
for i := range n+1 {
k := [3]int{x[i], y[i], z[i]}
cnt[k]++
}
// v個出現している値から二つを選ぶ選び方の個数
sameXYZ := 0
for _, v := range cnt {
sameXYZ += (v*(v-1))/2
}
よって式(3)を余事象で展開して整理すると、以下の式となります。
\begin{align}
&\sum_{l=0}^{N-1}\sum_{r=l+1}^{N}
[X(l) \neq X(r)][Y(l) \neq Y(r)][Z(l) \neq Z(r)]\\
=&\sum_{l<r}(1-[X(l)=X(r)])(1-[Y(l)=Y(r)])(1-[Z(l)=Z(r)]) \\
=&\sum_{l<r}1 \\
&- \sum_{l<r}[X(l)=X(r)]\\
&- \sum_{l<r}[Y(l)=Y(r)]\\
&- \sum_{l<r}[Z(l)=Z(r)]\\
&+ \sum_{l<r}[X(l)=X(r)][Y(l)=Y(r)]\\
&+ \sum_{l<r}[X(l)=X(r)][Z(l)=Z(r)]\\
&+ \sum_{l<r}[Y(l)=Y(r)][Z(l)=Z(r)]\\
&- \sum_{l<r}[X(l)=X(r)][Y(l)=Y(r)][Z(l)=Z(r)]
\end{align}
この式で表される各項はすべてO(N)以下で求められるため、制限時間以内に答えを出すことができました。
func samePairs[T comparable](x []T) int {
cnt := make(map[T]int)
for _, i := range x {
cnt[i]++
}
ans := 0
for _, x := range cnt {
ans += (x * (x - 1)) / 2
}
return ans
}
xy, xz, yz := make([][2]int, n+1), make([][2]int, n+1), make([][2]int, n+1)
xyz := make([][3]int, n+1)
for i := range n+1 {
xy[i] = [2]int{x[i], y[i]}
xz[i] = [2]int{x[i], z[i]}
yz[i] = [2]int{y[i], z[i]}
xyz[i] = [3]int{x[i], y[i], z[i]}
}
ans := n*(n+1)/2 -
samePairs(x) -
samePairs(y) -
samePairs(z) +
samePairs(xy) +
samePairs(xz) +
samePairs(yz) -
samePairs(xyz)
補足：計算の省略
最後の式で、$X(l)=X(r)$は区間[l,r)において$A(l,r)=B(l,r)$ であることを意味します。
同様に、$Y(l)=Y(r)$は$A(l,r)=C(l,r)$、$Z(l)=Z(r)$は$B(l,r)=C(l,r)$を意味します。
したがって、
- $X(l)=X(r)$ かつ $Y(l)=Y(r)$
- $X(l)=X(r)$ かつ $Z(l)=Z(r)$
- $Y(l)=Y(r)$ かつ $Z(l)=Z(r)$
はいずれも、結局 $A(l,r)=B(l,r)=C(l,r)$ と同値になります。
そのため、2要素の同時一致を数える3項と、3要素すべての同時一致を数える1項は、すべて同じ値になります。
よって、xy、xz、yz、xyz部分を圧縮することで、以下の式・コードとなります。
\begin{align}
&\sum_{l=0}^{N-1}\sum_{r=l+1}^{N}
[X(l) \neq X(r)][Y(l) \neq Y(r)][Z(l) \neq Z(r)]\\
=&\sum_{l<r}1 \\
&- \sum_{l<r}[X(l)=X(r)]\\
&- \sum_{l<r}[Y(l)=Y(r)]\\
&- \sum_{l<r}[Z(l)=Z(r)]\\
&+2* \sum_{l<r}[X(l)=X(r)][Y(l)=Y(r)][Z(l)=Z(r)]
\end{align}
xyz := make([][3]int, n+1)
for i := range n+1 {
xyz[i] = [3]int{x[i], y[i], z[i]}
}
ans := n*(n+1)/2 -
samePairs(x) -
samePairs(y) -
samePairs(z) +
2*samePairs(xyz)
以上です。