ABC290 E問題 Make it Palindrome

部分文字列

与えられた長さ $N$ の数列 $A$ の全ての連続部分列を対象としています。
部分文字列の最初の位置を $p$ 、最後の位置を $q$ とする部分文字列については、
$p と q$ , $p+1 と q-1$ , $p+2 と q-2, \cdots$ を見て、異なる数であれば変更する必要があります。
つまり、次のようなコードで答えが求められます。

ans = 0
for p = 1:N-1
  for q = p+1:N
    for k = 0:div(p-q,2)
      if A[p+k] != A[q-k]
        ans += 1
      end
    end
  end
end

しかし、これでは時間がかかりすぎます。

比較回数を減らす

ここで、 $A[i] と A[j]$ が異なるとします。 $i,j$ が回文の対称位置になる部分文字列は、
$p =i, q = j$
$p =i-1, q = j+1$
$p =i-2, q = j+2$
$\cdots$
で、 $p=i-k \ge 1$ かつ $q=j+k \le N$ を満たす 0 以上の $k$ の個数だけあります。
この個数は、 $min\left( i, N-j+1 \right)$ と書けます。
これを利用して、外側のループを $p,q$ ではなく $i,j$ で回すことにすると、
次のようなコードになります。

ans = 0
for i = 1:N-1
  for j = p+1:N
    if A[i] != A[j]
        ans += min( i, N-j+1 )
    end
  end
end

ところが、これでも時間がかかってしまいます。

二重ループの解消

時間がかかるのは $N$ の2重ループになっていて、 $O(N^2)$ の計算量になっているためです。
これを何とかしたいです。

そこで、 $i=1$ を固定して $j = 2,3,\cdots,N$ としてみます。このとき $min\left( i, N-j+1 \right)=1$ となるため
まとめて処理できそうです。
ところが $A[i] と A[j]$ をひとつずつ比較してしまうと2重ループのままです。
そこで、
$(異なる数の個数) = (全部の個数) - (同じ数の個数)$
という関係を利用します。
あらかじめ数列 $A$ に 1 が何個、2が何個、…を数えておく(個数の配列を用意する)と、
全部で( $i = 1$ も含めて) $N$ 個の数のうち、 $A[1]$ と同じ数が( $i = 1$ も含めて)何個あるかが
すぐわかります。この差から、 $j = 2,3,\cdots,N$ のうち $A[1]$ と異なる個数が求められます。

これで $i=1$ については数えられたので、個数の配列から $A[1]$ の分を引きます。

次に $i=2$ とすると、 $min\left( i, N-j+1 \right)=2$ になりそうですが、
$i=2, j=N$ では値が異なります。そこで、 $i=1$ の次は $j=N$ の場合について数えていきます。

$i=1$ で $j=2,\cdots,N$ の場合
$j=N$ で $i=2,\cdots,N-1$ の場合
$i=2$ で $j=3,\cdots,N$ の場合
$j=N-1$ で $i=3,\cdots,N-2$ の場合
$i=3$ で $j=4,\cdots,N-2$ の場合
$j=N-2$ で $i=4,\cdots,N-3$ の場合
$\cdots$
のように $i,j$ 交互に処理を進めていきます。

コード例 (Julia)

function solve()

  # 入力
  readInt()  =  parse(Int,readline())
  readInts() =  parse.(Int,split(readline()))

  # 入力データ
  N = readInt()
  A = readInts()

  # 各数字が何回現れるか
  C = zeros(Int,N)
  for i = 1:N
    C[A[i]] += 1
  end

  # A[i] と A[j] (i<j) が異なるの場合、A[i] とA[j] が回文で対応する位置になるような
  # 部分文字列は min(i,N-j+1) 種類あり、この値が答えに加算される。

  head = 1
  tail = N
  ans = 0

  sumC = N

  while head < tail
    # head から tail の間に、A[head] と異なる値がいくつあるか
    p = sumC - C[A[head]]
    ans += p * head

    # 処理済みのA[head]の分を減らす
    C[A[head]] -= 1
    sumC -= 1

    # head+1 から tail の間に、A[tail] と異なる値がいくつあるか
    p = sumC - C[A[tail]]
    ans += p * head

    # 処理済みのA[tail]の分を減らす
    C[A[tail]] -= 1
    sumC -= 1

    # head, tail を内側に 1つずらす
    head += 1
    tail -= 1
  end

  println(ans)

end  # function solve

# main