atcoder ABC149E で時間超過になってしまいます。

Question

atcoder の問題を解いているのですが、正解にたどり着けません。
自分で考えてもわからなかったので、コードで誤っている部分を指摘してもらえるとうれしいです。

atcoder ABC 149E
問題URL

問題文

高橋 君は、あるパーティに特別ゲストとしてやってきました。 そこには一般ゲストが N 人おり、一般ゲスト i (1≤i≤N) の パワー は Ai です。高橋君は 握手 を M 回行うことで、パーティ全体の 幸福度 を上げることにしました (握手開始前のa幸福を 0 とします)。 握手は、以下の手順によって行われます。

高橋君が左手で手を握る (一般) ゲスト x と右手で手を握るゲスト y を決める (両手で同じゲストの手を握っても良い)。高橋君が実際にこれら二本の手を握ることで、幸福度が Ax+Ay 上がる。
ただし、全く同じ握手を二回以上行ってはいけません。厳密には、次の条件を満たす必要があります。k 回目の握手で、左手でゲスト xk と、右手でゲスト yk と手を握ったとする。このとき、 (xp,yp)=(xq,yq) を満たすような p,q(1≤p<q≤M) が存在しない。M 回の握手を行った後、最終的な幸福度は最大でいくらにできるでしょうか。

制約

1≤N≤10^5
1≤M≤N^2
1≤Ai≤10^5
入力は全て整数である。

自分のコード

#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;

int N, M, done = 0;
vector<int> A, B;
long long ans = 0;

bool judge(int mid){
    int count = 0;
    for(int i=0; i<N; i++){
        int j = 0;
        while(mid <= A[i]+A[j] && j<N) j++;
        count += j;
    }
    if(count >= M) return true;
    else return false;
}

int main(){
    cin >> N >> M;
    A.resize(N);
    B.resize(N);
    for(int i=0; i<N; i++) cin >> A[i];
    sort(A.rbegin(), A.rend());
    for(int i=1; i<N+1; i++){
        B[i] = B[i-1] + A[i-1];
    }

    int right = A[0]*2 + 1;
    int left = A[N-1]*2;
    int mid;
    while(left+1 < right){
        mid = (left + right)/2;
        if(judge(mid)){
            left = mid;
        }
        else right = mid;
    }

    for(int i=0; i<N; i++){
        if(done == M) continue;
        if(left > A[i] + A[0]) continue;
        int j = 0;
        while(left <= A[i] + A[j] && j<N) j++;
        ans += B[j] + A[i]*j;
        done += j;
    }
    cout << ans << endl;
}

参考にしたコード

#include <iostream>
#include <algorithm>
#include <string>
#include <vector>
#include <cmath>
#include <map>
#include <queue>
#include <iomanip>
#include <set>
#include <tuple>
#define mkp make_pair
#define mkt make_tuple
#define rep(i,n) for(int i = 0; i < (n); ++i)
using namespace std;
typedef long long ll;
const ll MOD=1e9+7;
 
ll N,M;
vector<ll> A;
 
vector<ll> sum;
 
bool judge(ll lim){
  ll res=0;
  int itr=0;
  for(int i=N-1;i>=0;i--){
    while(itr<N&&A[i]+A[itr]>=lim){
      itr++;
    }
    res+=itr;
  }
 
  if(res>=M) return true;
  else return false;
}
 
int main(){
    cin>>N>>M;
    A.resize(N);
    rep(i,N) cin>>A[i];

    sort(A.rbegin(),A.rend());

    sum.resize(N+1,0);
    rep(i,N) sum[i+1]=sum[i]+A[i];

    ll ok=0,ng=(A[0]*2+1)*M;
    while(abs(ok-ng)>1){
        ll mid=(ok+ng)/2;
        if(judge(mid)) ok=mid;
        else ng=mid;
    }

    ll ans=0;
    ll itr=0;
    ll num=0;
    for(int i=N-1;i>=0;i--){
        while(itr<N && A[i]+A[itr]>=ok) itr++;
        ans+=A[i]*itr+sum[itr];
        num+=itr;
    }
    ans-=ok*(num-M);
    cout<<ans<<endl;

    return 0;
}

Answer 1

bool judge(int mid){
    int count = 0;
    for(int i=0; i<N; i++){
        int j = 0;
        while(mid <= A[i]+A[j] && j<N) j++;
        count += j;
    }
    if(count >= M) return true;
    else return false;
}

ここがO(N^2)

    for(int i=0; i<N; i++){
        if(done == M) continue;
        if(left > A[i] + A[0]) continue;
        int j = 0;
        while(left <= A[i] + A[j] && j<N) j++;
        ans += B[j] + A[i]*j;
        done += j;
    }

ここがO(M)になるので時間制限上厳しいでしょう。

そのほかにも配列の範囲外へのアクセスがあるのと、M回を超えて握手する可能性があります(done > M)。テストケースによってはWAになるでしょう。

---

参考コードとの比較

  int itr=0;
  for(int i=N-1;i>=0;i--){
    while(itr<N&&A[i]+A[itr]>=lim){
      itr++;
    }
    res+=itr;
  }

こちらのコードの場合、itrの初期化位置がforループの外なので、次のiの計算でも前のitrがそのまま残った状態です。itrは増える一方で、最大Nまでインクリメントされるだけなので、whileループの中は高々N回しか呼ばれません。結果計算量はO(N)になります。

    ll itr=0;
    ll num=0;
    for(int i=N-1;i>=0;i--){
        while(itr<N && A[i]+A[itr]>=ok) itr++;
        ans+=A[i]*itr+sum[itr];
        num+=itr;
    }

こちらも同様にitrがインクリメントされるのは高々N回なのでO(N)です。

Comments

[goro_gnm]

たしかに計算量的にだめなのかなと思いましたが、参考にさせてもらった他の人のコードでは計算量は同じに見えるのですがACになっていました。（質問文に追記しましたので見てもらえるとありがたいです。）
違いは何なんでしょうか。

M回を超えてしまう場合について考えられていませんでした。ありがとうございます。
配列の範囲外へのアクセスについてはもう一度見直してみます。

[goro_gnm]

よくわかりました、ありがとうございます！