各行動のタイミングにおいて、高橋君の現在の座標 \(X\) から最も近いクッキーとして考えられるのは「座標が \(X\) 未満のうち最大のもの」「座標が \(X\) より大きいうち最小のもの」のどちらかです。よってこの2つを高速に求めることができればこの問題に答えることができます。

解法1 ordered set

C++ などには set というデータ構造があり、以下の操作を全て \(O(\log N)\) で行うことができます

• set に要素を1つ追加する
• setから要素を1つ削除する
• setに含まれる要素のうち X 未満も最大値を求める
• setに含まれる要素のうちXより大きな最小値を求める

よって、クッキーのある座標をsetで管理しながら高橋君の行動をシミュレーションすることで \(O(N\log N)\) でこの問題を解くことができます。

writer解(C++)

#include<bits/stdc++.h>
using namespace std;

int main(){
	int n;
	cin >> n;
	set<int>s;
	for(int i=0;i<n;i++){
		int ai;
		cin >> ai;
		s.insert(ai);
	}
	
	long long ans=0;
	int pos=0;
	for(int i=0;i<n;i++){
		auto it=s.lower_bound(pos);
		int nxt;
		if(it==s.begin()){
			nxt=*it;
		}else if(it==s.end()){
			nxt=*--it;
		}else{
			int cand1=*it;
			int cand2=*--it;
			if(abs(cand1-pos)<abs(cand2-pos)){
				nxt=cand1;
			}else{
				nxt=cand2;
			}
		}
		ans+=abs(nxt-pos);
		s.erase(nxt);
		pos=nxt;
	}
	cout << ans << endl;
}

解法2 2-pointers

クッキーのある座標をソートしておき、「座標が負のものを、原点に近い順に取り出せるキュー」「座標が正のものを、原点に近い順に取り出せるキュー」を用意します。

どのタイミングでも、高橋君が次に拾うクッキーは、2つのキューの先頭のどちらかになります。よって、クッキーのある座標を2つのキューで管理しながら高橋君の行動をシミュレーションすることでソートの後は \(O(N)\) でこの問題を解くことができます。

以下の実装例では明示的にキューを用意せず、座標をソートしたリストと、次に取り出す要素のindexを表す変数を用いて同じ仕組みを実現しています。

writer解(Python)

N=int(input())
A=list(map(int,input().split()))

N+=1
A.append(0)
A.sort()
P=A.index(0)
L=P-1
R=P+1
ans=0

pos=0
for _ in range(N-1):
  if L==-1:
    ans+=A[R]-pos
    pos=A[R]
    R+=1
  elif R==N:
    ans+=pos-A[L]
    pos=A[L]
    L-=1
  else:
    if pos-A[L]<=A[R]-pos:
      ans+=pos-A[L]
      pos=A[L]
      L-=1
    else:
      ans+=A[R]-pos
      pos=A[R]
      R+=1

print(ans)