Official

F - Double Chance Editorial by en_translator


For a fixed \(K\), there are \(K\) choices for each draw, for a total of \(K^2\) choices, so the desired value is \( \left(\displaystyle\sum_{i=1}^{K}\sum_{j=1}^{K}\max(A_i,A_j)\right)/K^2 \). Since \(K^2\) can be computed easily, we consider how to find \(\displaystyle\sum_{i=1}^{K}\sum_{j=1}^{K}\max(A_i,A_j)\). Since it is troublesome to compute for every \(K\), we will try to reuse the value for \(K-1\). Considering the difference, we have

\[ \displaystyle\sum_{i=1}^{K}\sum_{j=1}^{K}\max(A_i,A_j) =\displaystyle\sum_{i=1}^{K-1}\sum_{j=1}^{K-1}\max(A_i,A_j) +2\displaystyle\sum_{i=1}^{K-1}\max(A_i,A_K) +A_K, \]

where \(\displaystyle\sum_{i=1}^{K}\sum_{j=1}^{K}\max(A_i,A_j)=0\) for \(K=0\). All that left is to find \(\displaystyle\sum_{i=1}^{K-1}\max(A_i,A_K)\) fast enough.
The value \(\displaystyle\sum_{i=1}^{K-1}\max(A_i,A_K)\) can be decomposed into two parts by \(1\leq i\leq K-1\), depending on which of \(A_i\) and \(A_K\) is larger:

\[ \displaystyle\sum_{i=1}^{K-1}\max(A_i,A_K) =\displaystyle\sum_{\substack{1\leq i\leq K-1 \\A_i\leq A_K}}A_K +\displaystyle\sum_{\substack{1\leq i\leq K-1 \\A_i>A_K}}A_i. \]

Therefore, it is sufficient to find the number of \(i\)’s such that \(A_i \leq A_K\) and the sum of \(A_i\)’s such that \(A_i>A_K\); there are several approaches. We introduce two of them.

Solution 1. finding for each \(K=1,2,\ldots\)

Since each \(A_i\) is at most \(2\times 10^5\), we can use two Fenwick Trees.

Let \(F_1[i]=(\text{The number of }j\text{'s such that }(1\leq j\leq K-1\text{ and })A_j=i)\) and
\(F_1[i]=(\text{The number of }j\text{'s such that }(1\leq j\leq K-1\text{ and })A_j=i)\times i\),
then we can find the sought values by \(\displaystyle\sum_{\substack{1\leq i\leq K-1 \\A_i\leq A_K}}A_K=(F_1[1]+F_1[2]+\cdots+F_1[A_K])\times A_K\) and
\(\displaystyle\sum_{\substack{1\leq i\leq K-1 \\A_i>A_K}}A_i=F_2[A_K+1]+F_2[A_K+2]+\cdots+F_2[M]\),
where \(M=2\times 10^5\) is the upper bound of \(A_i\). All that left is to do the following for each \(K=1,2,\ldots,N\):

  1. Find \(\displaystyle\sum_{i=1}^{K-1}\max(A_i,A_K)\)
  2. Find \(\displaystyle\sum_{i=1}^{K}\sum_{j=1}^{K}\max(A_i,A_j)/K^2\) using the result of 1
  3. Update \(F_1\) and \(F_2\) (by adding \(1\) to \(F_1[A_K]\) and \(A_K\) to \(F_2[A_K]\))

(2. and 3. are interchangeable.) Then you can find the answer. The time complexity is \(O(N\log\max(A_i))\), which is fast enough.

Note that, if \(A_i\) can be up to \(10^9\), you cannot reserve the memory for the Fenwick trees in this approach, so you need to do coordinate compression etc..

Solution 2. finding in descending order of \(A_K\)

Alternatively, the number of \(i\)’s such that \(A_i \leq A_K\) and the sum of \(A_i\) such that \(A_i>A_K\) can be found with the following two Fenwick trees by sorting \((A_K,K)\) in descending order beforehand:

\(F_1[i]=(\text{The number of }i\text{'s such that }(A_i,i)>(A_K,K))\) (That is, \(1\) if \((A_i,i)>(A_K,K)\) and \(0\) otherwise);
\(F_1[i]=(\text{The number of }i\text{'s such that }(A_i,i)>(A_K,K))\times A_i\) (That is, \(A_i\) if \((A_i,i)>(A_K,K)\) and \(0\) otherwise);

(where integer pairs are compared lexicographically) Then you can find the sought values by \(\displaystyle\sum_{\substack{1\leq i\leq K-1 \\(A_i,i)\leq (A_K,K)}}A_K=\{(K-1)-(F_1[1]+F_1[2]+\cdots+F_1[K-1])\}\times A_K\) and
\(\displaystyle\sum_{\substack{1\leq i\leq K-1 \\(A_i,i)>(A_K,K)}}A_i=F_2[1]+F_2[2]+\cdots+F_2[K-1]\).
Here note that, if \(A_i=A_K\), it can be included in either sum as long as there is included twice. All that left is to compute the value \(\displaystyle\sum_{i=1}^{K-1}\max(A_i,A_K)\) and update \(F_1, F_2\) in descending order of \(A_K\).

The time complexity for this solution is \(O(N\log N)\). As for this approach, you can apply it even if \(A_i\) can be up to \(10^9\). By the way, in our case, as \(\displaystyle\sum_{i=1}^N A_i\leq 4\times 10^{10}<10^{12}\), we can use only one Fenwick tree by letting \(F[i]=F_1[i]\times 10^{12}+F_2[i]\) to improve the constant factor, although it is rarely required.

Sample code in C++ (Approach \(1\)):

#include <bits/stdc++.h>
#include <atcoder/fenwicktree>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
#define M 200000


int main() {
  int n,x;
  mint d,s,ans;

  cin>>n;
  fenwick_tree<mint> fw1(M+1);
  fenwick_tree<mint> fw2(M+1);
  s=0;
  for(int i=0;i<n;i++){
    cin>>x;
    d=fw1.sum(0,x+1)*mint(x)+fw2.sum(x+1,M+1);
    s=s+mint(2)*d+mint(x);
    cout<< (s*(mint(i+1).inv().pow(2))).val() <<"\n";
    fw1.add(x,mint(1));
    fw2.add(x,mint(x));
  }
  return 0;
}

Sample code in C++ (Approach \(2\)):

#include <bits/stdc++.h>
#include <atcoder/fenwicktree>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
#define C (long long)1e+12



int main() {
  int n,k;
  mint s;
  long long x;

  cin>>n;

  vector<int>a(n);
  vector<pair<int,int> >ord(n);
  fenwick_tree<long long> fw(n);
  vector<mint>d(n);

  for(int i=0;i<n;i++){
    cin>>a[i];
    ord.push_back({a[i],i});
  }
  sort(ord.begin(),ord.end(),greater<pair<int,int> >());

  for(int i=0;i<n;i++){
    k=ord[i].second;
    x=fw.sum(0,k);
    d[k]=mint((k-(x/C))*a[k]+(x%C));
    fw.add(k,C+a[k]);
  }

  s=0;
  for(int i=0;i<n;i++){
    s+=mint(2)*d[i]+mint(a[i]);
    cout <<(s*(mint(i+1).inv().pow(2))).val()<<"\n";
  }
  return 0;
}

posted:
last update: