Official

H - Count Arithmetic Subsequences Editorial by en_translator


Every length-\(1\) subsequence is an arithmetic sequence; there are \(N\) of them. Putting them aside, we will count arithmetic sequences of length \(2\) or greater.

\(O(N^4)\) solution

Consider DP (Dynamic Programming) with the following states:

\(\mathrm{dp}[i][j][l] \coloneqq\) the number of arithmetic sequences whose initial and second terms are \(A_i\) and \(A_j\), and length is \(l\,(l \geq 2)\).

Iterate over \(A_i\) in descending order of \(i\). For each fixed \(A_i\), brute-force over the second term \(A_j\,(i<j)\) and the length \(l\). If \(l=2\), then \((A_i,A_j)\) is an arithmetic sequence of length \(2\), so add \(1\) to \(\mathrm{dp}[i][j][2]\). If \(l\geq 3\), further exhaustively search over \(A_k\,(j<k)\). When \(A_k-A_j=A_j-A_i\), an arithmetic sequence with initial term \(A_i\), second term \(A_j\), and length \(l\), can be obtained by prepending \(A_i\) to an arithmetic sequence with initial term \(A_i\), second term \(A_k\), and length \((l-1)\). Thus, add \(\mathrm{dp}[j][k][l-1]\) to \(\mathrm{dp}[i][j][l]\).

With \(O(N^3)\) states and \(O(N)\) transition cost, the overall complexity is \(O(N^4)\); this runs fast enough under the constraints of this problem.

Sample code (PyPy)

\(O(N^3)\) solution

The solution above can be further optimized. Consider DP with the following states:

\(\mathrm{dp}[i][j][l] \coloneqq\) the number of arithmetic sequences whose initial term is \(A_i\), length is \(l\,(l \geq 2)\), and common difference is \(d\).

Iterate over \(A_i\) in descending order of \(i\). For each fixed \(A_i\), brute-force over the second term \(A_j\,(i<j)\) and the length \(l\). Here, the common difference is determined as \(d\coloneqq A_j-A_i\). If \(l=2\), then \((A_i,A_j)\) forms a length-\(2\) arithmetic sequence, so add \(1\) to \(\mathrm{dp}[i][2][d]\). If \(l\geq 3\), then an arithmetic sequence with initial term \(A_i\), length \(l\), and common difference \(d\), can be obtained by prepending \(A_i\) to an arithmetic sequence with initial term \(A_j\), length \((l-1)\), and common difference \(A_i\). Therefore, add \(\mathrm{dp}[j][l-1][d]\) to \(\mathrm{dp}[i][l][d]\).

Since \(d\) can take as large value as \(10^9\), indices for \(d\) can be stored as an associative array. If you use a map in C++, the time complexity is \(O(N^3\log N)\).

Alternatively, one can enumerate all possible public difference for each \(i\) and apply coordinate complexity to store indices for \(d\) as an array. In this case, the time complexity is \(O(N^3)\).

Sample code (PyPy, Coordinate compression)

posted:
last update: