Official

I - Double Sum 3 Editorial by en_translator


When the blackboard has some integers, it is optimal to pick \(l\) and \(r\) to maximize \(r-l\) in order to minimize the number of operations. That is, we can pick \((l,r)\) such that:

  • \(l-1\) is not written on the blackboard,
  • all integers from \(l\) through \(r\) are written on the blackboard, and
  • \(r+1\) is not written on the blackboard,

and removing them.

We will solve the problem based on this fact.

Fixing an integer \(1\le c\le N\), consider the number \(g(c)\) of pairs \((L,R)\) such that the procedure corresponding to \(f(L,R)\) requires an operation with \(l=c\). Then the sought answer is \(\displaystyle \sum_{c=1}^Ng(c)\). From now on, we will consider how to evaluate \(g(c)\) for each \(c\).

According to the observation above, we need to pick \(l=c\) if and only if \(S=\{A_L,A_{L+1},\ldots,A_R\}\) satisfies:

  • \(c-1\notin S\) and
  • \(c\in S\).

Thus, \(g(c)\) coincides with the number of pairs \((L,R)\) with \(c-1\notin S\) subtracted by the number of pairs \((L,R)\) with \(c-1\notin S\) and \(c\notin S\).

Here, let \(G_c\) be the set of \(i\) with \(A_i=c\). Then \(c\notin S\) if and only if any integer from \(L\) through \(R\) does not belong to \(G_c\). The number of such pairs \((L,R)\) is \(\displaystyle \sum_{k=0}^{|G_c|}\frac{(X_{k+1}-X_k)(X_{k+1}-X_k-1)}2\), where \(G_c=\{X_1,X_2,\ldots,X_{|G_c|}\}\) \((X_1<X_2<\ldots<X_{|G_c|}),\) \(X_0=0\), and \(X_{|G_c|+1}=N+1\). The number of pairs with \(c-1\notin S\) and \(c\notin S\) can be found in a similar way.

By this computation, we can find \(g(c)\) for each \(g(c)\), and the answer can be found as their sum. The complexity is \(O(N)\).

The following implementation costs \(O(N\log N)\) for readability.

Sample code (Python3)

n = int(input())
a = list(map(int, input().split()))
g = [[] for i in range(n + 1)]
for i in range(n):
    g[a[i] - 1].append(i)

def f(x):
    res = 0
    for i in range(1, len(x)):
        c = x[i] - x[i - 1]
        res += c * (c - 1) // 2
    return res

ans = 0
for i in range(1, n + 1):
    ans += f([-1] + g[i] + [n])
    ans -= f([-1] + sorted(g[i - 1] + g[i]) + [n])
print(ans)

posted:
last update: