#問題概要
要素数$N$の配列$A$が与えられる。
数直線上に配列$A$の各要素が並んでいる。原点からスタートして$A(i)$への移動を繰り返し、再び原点に戻ってくる。この時、点$a$から点$b$への移動は$|a-b|$円かかる。
$i = 1, 2, ..., N$について点$i$への移動をとりやめた時にかかる金額の総和をそれぞれ求めよ。
#制約
- $2 \leqq N \leqq 10^5$
- $−5000 \leqq A(i) \leqq 5000(1 \leqq i \leqq N)$
- 入力値はすべて整数である。
#考えたこと
まず配列$A$から要素$A(i)$を抜いた配列を作り総和を取っていくのを$N$回繰り返すというやり方が思いつくが、制約条件で$N$が最大$10^5$なので$O(N^2)$だと計算量が$10^{10}$となり、かなり厳しい。
なので$for$文一つにして計算量を$O(N)$に収めたい。
ここでまず配列Aの要素間の差分の総和をとってうまく引き算したらよさそうと考えて、入力例で実験した。
入力例1
3
3 5 -1
すると以下のように点の移動と点から点への移動コストが整理される。
この図から配列Aの両端に0を追加し、 $B[i]=|A[i]-A[i+1]|$ となる配列Bを作り、配列$B$の総和から、 $B[i]+B[i+1]$ を引き、 $|A[i]-A[i+2]|$ を足せば求める値が出せることがわかる。
例えば、上記の入力例で言うと、移動のとりやめを考慮しない時の総和は12。
点$i=0$の時、$B[0]=3$と$B[1]=2$を$B$の総和から引いて、$|A[0]-A[2]|=5$を$B$の総和に足す。
あとは出力すればおーけーです。これをN回繰り返すだけ。
ロジックを簡潔に示すと以下のようになる。
a.sum()
- (abs(a[i] - a[i - 1]) + abs(a[i + 1] - a[i])) //=b[i] + b[i+1]
+ abs(a[i + 1] - a[i - 1])
#解答
using System;
using System.Linq;
using System.Collections.Generic;
class Program
{
static void Main(string[] args) {
int n = int.Parse(Console.ReadLine());
List<int> a = new List<int>();
a.Add(0);
a.AddRange(Console.ReadLine().Split().Select(int.Parse).ToArray());
a.Add(0);
List<int> b = new List<int>();
for (int i = 0; i < n+1; i ++) b.Add(Math.Abs(a[i]-a[i+1]));
int sum = b.Sum();
for (int i = 0; i < n; i ++) {
int result = sum;
result -= b[i]+b[i+1];
result += Math.Abs(a[i]-a[i+2]);
Console.WriteLine(result);
}
}
}
#犯したミス
ACするまでの過程で犯したミスを書いておきます。
- 上記の解法を実際のコードに落としていく時にループごとに
result
を初期化するのを失念していた。 - 1.の解消のためにbの総和を取るのをfor文中に入れたらTLEした。以下の記事を参照すると、計算時間は for < foreach < Sumとなっている。よって今回の問題ではfor文中にsumを入れると$O(N^2)$となるので避けるべき。
これぐらいですかね。他はインデックスに気をつけるとかです。
あと僕はC#で書いてるので、配列をListにして扱うほうが動的な処理が可能なのでよいかと思います。