$O(N \log M)$ ですが実装が楽な解法を紹介します。
公式解説にて $dp[i-1][A_i]$ を除いた $dp[i-1]$ の最大値 $mx$ が知りたい場面があります。これを $O(\log M)$ で求めます。

初めに、求める値 \(mx\) は \(\max \left( \max \left( dp[i-1][1] , \cdots ,dp[i-1][A_i-1] \right) , \max \left( dp[i-1][A_i+1] , \cdots ,dp[i-1][M] \right) \right)\) と言い換えられます。つまり、配列の区間内の最大値を求める処理(RMQやrange max queryなどと言われます)を \(2\) 回行えばよいです。

ここでノード数が \(M\) のsegmenttreeを用意し、 \(j\) 番目のノードに \(dp[i-1][j]\) を乗せると \(mx\) はRMQを \(2\) 回解けば求まります。よって \(O(\log M)\) で \(mx\) が求まります。

また、 \(A_i \neq j\) ならば \(dp[i-1][j] = dp[i][j]\) であるので、 \(mx\) を介して求めた \(dp[i][A_i]\) でsegmenttreeの \(A_i\) 番目のノードを更新させれば \(dp[i+1][A_{i+1}]\) に対応する \(mx\) を求めるのに必要となるものが揃います。