ABC379 C問題の解説（Python）

はじめに

この記事では、ABC379のC問題を解説していきます。

ABC（AtCoder Beginner Contest）とは、AtCoderが開催している、競技プログラミングコンテストです。

ABC379　C – Sowing Stones

問題

1 から N まで番号がつけられた N 個のマスが一列に並んでいます。最初、 M 個のマスに石が入っており、マス \(X_i\) には \(A_i\) 個の石が入っています。

以下の操作を好きな回数行うことができます。

マス \(i\) に石があるとき、マス \(i\) から石を 1 つマス \(i+1\) に移動させる。

N 個のマスすべてに石がちょうど 1 個ずつ入っている状態にするために必要な操作回数の最小値を求めてください。ただし、不可能な場合は -1 を出力します。

制約

\(2 \leqq N \leqq 2×10^9\)
\(1 \leqq M \leqq 2×10^5\)
\(M \leqq N\)
\(i ≠ j\) のとき、\(X_i ≠ X_j\)
\(1 \leqq A_i \leqq 2×10^9\)

思考の筋道

左端から実際にシミュレーションして考えてみましょう。

方針の検討

前提として、各マスの石の合計が N 個でなければなりません。

また、各マスに石は1個だけなので、あるマスに石を1個配置すれば、残りの石はすべて右のマスに流すことになります。

ただし、N は最大で \(2×10^9\) なので、N に関するループは行えません。すべてのマスを調べられないとなると、どうすればよいでしょうか？

すべてのマスではなく、石が置かれているマスについてのみ考えることにしましょう。

ある石が置かれているマスAから、次の石が置かれているマスBまで、どのように石を運ぶか考えます。マスAと途中のマスに石を1個ずつ配置し、残った石はすべてマスBにマージします。なお、途中のマスに石を1個ずつ配置できないときは、-1を出力しましょう。

マスAからマスBに移動時の操作回数と石の数の変化

マスAからマスBまでの距離を \(d\) とします。仮にマスAにあるすべての石をマスBに移動する場合、必要な操作回数は（マスAにある石の個数）× \(d\) です。

しかし実際は、マスAと途中のマスに石を1個ずつ配置していきますので、その分だけ操作回数は減ります。減る操作回数は、\(1+2+\cdots+d = d ×(d+1)÷2\)です。

これで、マスAからマスBに石を移動するのに必要な操作回数が求められました。

次に、石の数ですが、マスAからマスBの距離分だけ減少した後、マスBに置かれている石の数だけ増加します。

操作回数と石の数について調べられましたので、あとはコーディングをしていきましょう。

残る注意点

入力で与えられる X は昇順ソートされていません。X と A を、Xの昇順でソートする必要があります。（M は最大で \(2×10^5\) よりソート可能）
石が置かれているマスについてのみ考えると書きましたが、最後の右端のマスも考える必要があります。下のコードでは、Xのソート後に、リストの末尾にマスNを石の個数0で追加して調べました。

なかなか難しい問題でしたね・・・。

コード

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

if sum(A) != N:
  print(-1)
  exit()

XA = sorted(zip(X, A))
XA.append((N, 0))

pos = 1
stock = 0
ans = 0

for row in XA:
  d = row[0] - pos
  if d > stock:
    print(-1)
    exit()
  else:
    ans += d * stock - d * (d + 1) // 2
    stock += row[1] - d
    pos = row[0]

print(ans)

以上、ABC379のC問題の解説でした！

では、またね。

リンク

ABC379

情報科学の部屋へ

ABC379 C問題の解説（Python）

はじめに