결국 가장 먼 곳까지 가야 한다는 규칙을 발견하고, 이를 이용하는 것이 핵심인 문제입니다.
가장 먼 곳까지 가서 책을 꽂고 오면 항상 가장 먼 곳이 이전보다 가까워집니다.
이전보다 가까워지는 거리를 최대로 하기 위해서는 가장 먼 곳에서 먼 순서대로 $M$개의 책을 꽂고 와야 합니다.
이 때의 비용은 가장 먼 곳을 한 번 가는 것과 같으므로 최선의 선택이 됩니다.
위의 알고리즘을 구현하기 위해서는 우선 책의 위치를 정렬하고 시작점인 0의 위치를 찾은 뒤 음수 쪽에서 $M$개씩 책을 가져오고, 0에 다시 도착하면 양수 쪽에서 $N$개씩 책을 가져오는 것을 반복해야 합니다.
$M=1$일 때 모든 원소를 한 번씩 방문하게 되므로 $O\left(N\right)$의 시간이 걸리며, 위치 순으로 정렬할 때 $O\left(N\log N\right)$ 시간으로 총 $O\left(N\log N\right)$ 시간이 걸립니다.
코드
fn solve(pos: &mut [i32], n_lift: usize) -> u32 {
pos.sort_unstable();
let zero = pos.partition_point(|&p| p < 0);
let mut left = 0;
let mut walk = 0;
while left < zero {
walk += -pos[left] << 1;
left += n_lift;
}
let mut right = pos.len() - 1;
if zero <= right {
loop {
walk += pos[right] << 1;
if right < zero + n_lift {
break;
} else {
right -= n_lift;
}
}
}
walk -= pos[0].abs().max(pos[pos.len() - 1]).abs();
walk as u32
}