Zenn
🍡

[鉄則A09] Winter in ALGO Kingdom

2024/04/27に公開
Ruby
AtCoder
競プロ
累積和
いもす法
tech

https://atcoder.jp/contests/tessoku-book/tasks/tessoku_book_i

差分＋二次元の累積和問題。

問題を要約する

  • 矩形範囲に雪が降る
  • 日ごとの範囲データがある
  • 最後に積雪状態はどうなっている？

解き方

  1. 矩形の隅の四点を±1する
  2. それを矩形の個数分繰り返す
  3. 最後にまとめて1回だけ累積和テーブルを作る

例1でシミュレーション

5 x 5 のマップ内の (x1, y1) - (x2, y2) の範囲に雪が降るとしたとき、一つ目は (0, 0) - (2, 2) なので、矩形の隅の四点±1する。

  • 左上(角の内) + 1
  • 右下(角の外) + 1
  • 左下(辺の外) - 1
  • 右上(辺の外) - 1

すると、

0 1 2 3 4 5
0 1 0 0 -1 0 0
1 0 0 0 0 0 0
2 0 0 0 0 0 0
3 -1 0 0 1 0 0
4 0 0 0 0 0 0
5 0 0 0 0 0 0

になる。

このとき、仮に (x2, x2) が (4, 4) だとすれば右下(角の外)が W x H を超えるので右と下にパディングが必要になる。これをせずにかなりはまった。ここでは一長一短あるものの範囲チェックして W x H を超えないようにする方法もある。

続いて二つ目の (1, 1) - (3, 3) も±1すると、

0 1 2 3 4 5
0 1 0 0 -1 0 0
1 0 1 0 0 -1 0
2 0 0 0 0 0 0
3 -1 0 0 1 0 0
4 0 -1 0 0 1 0
5 0 0 0 0 0 0

になる。

それを→↓の順で累積和にする。

0 1 2 3 4 5
0 0 0 0 0 0 0
1 0 1 1 1 0 0
2 0 1 2 2 1 0
3 0 1 2 2 1 0
4 0 0 1 1 1 0
5 0 0 0 0 0 0

それの一行目と左端の列を外したのが答えになる。

1 2 3 4 5
1 1 1 1 0 0
2 1 2 2 1 0
3 1 2 2 1 0
4 0 1 1 1 0
5 0 0 0 0 0

あとで気づいたが、最後に範囲の合計を出せとは言われていないため、累積和テーブルの一行目と左端の一列はなくてもいい。いや、むしろない方が実装が楽になる。

単純な実装 (TLE)

H, W, N = gets.split.collect(&:to_i)                # => [5, 5, 2]
R = N.times.collect { gets.split.collect(&:to_i) }  # => [[1, 1, 3, 3], [2, 2, 4, 4]]
R = R.collect { |e| e.collect(&:pred) }             # => [[0, 0, 2, 2], [1, 1, 3, 3]]
S = Array.new(W) { Array.new(H, 0) }
R.each do |y1, x1, y2, x2|
  (x1..x2).each do |x|
    (y1..y2).each do |y|
      S[x][y] += 1
    end
  end
end
S                                                   # => [[1, 1, 1, 0, 0], [1, 2, 2, 1, 0], [1, 2, 2, 1, 0], [0, 1, 1, 1, 0], [0, 0, 0, 0, 0]]
puts H.times.collect { |y| W.times.collect { |x| S[y][x] } * " " }

わかりやすいが最悪 O(NWH) かかるので TLE になる。

解説の模範解答風 (AC)

1140 ms
H, W, N = gets.split.collect(&:to_i)                # => [5, 5, 2]
R = N.times.collect { gets.split.collect(&:to_i) }  # => [[1, 1, 3, 3], [2, 2, 4, 4]]

# 必ず 0-based で統一する
R = R.collect { |e| e.collect(&:pred) }             # => [[0, 0, 2, 2], [1, 1, 3, 3]]

# 範囲に対応する四角形の隅の四点を±1する
F = Array.new(H.next) { Array.new(W.next, 0) }
R.each do |y1, x1, y2, x2|
  F[y1 + 0][x1 + 0] += 1 # 左上(角の内)
  F[y2 + 1][x2 + 1] += 1 # 右下(角の外)
  F[y2 + 1][x1 + 0] -= 1 # 左下(辺の外)
  F[y1 + 0][x2 + 1] -= 1 # 右上(辺の外)
end

# 上と左に一行一列入れる
S = Array.new(H.next) { Array.new(W.next, 0) }

# 横方向の累積和 (S[v] = 前 + F)
(1..H).each do |y|
  (1..W).each do |x|
    S[y][x] = S[y][x.pred] + F[y.pred][x.pred]
  end
end

# 横方向の累積和 (S[v] += 上)
(1..W).each do |x|
  (1..H).each do |y|
    S[y][x] += S[y.pred][x]
  end
end

puts (1..H).collect { |y| (1..W).collect { |x| S[y][x] } * " " }

リファクタリング後 (AC)

こちらの実装が参考になった。

  • 二次元配列は2つもいらない
  • 累積和テーブルの一行目と左端の一列はいらない
  • 累積和は前の値を保持する変数を持っておけばよい

ということで、それを踏まえて改善したもの。

1098 ms
H, W, N = gets.split.collect(&:to_i)                # => [5, 5, 2]
R = N.times.collect { gets.split.collect(&:to_i) }  # => [[1, 1, 3, 3], [2, 2, 4, 4]]
R = R.collect { |e| e.collect(&:pred) }             # => [[0, 0, 2, 2], [1, 1, 3, 3]]

S = Array.new(H.next) { Array.new(W.next, 0) }
R.each do |y1, x1, y2, x2|
  S[y1 + 0][x1 + 0] += 1 # 左上(角の内)
  S[y2 + 1][x2 + 1] += 1 # 右下(角の外)
  S[y2 + 1][x1 + 0] -= 1 # 左下(辺の外)
  S[y1 + 0][x2 + 1] -= 1 # 右上(辺の外)
end

H.times do |y, a = 0|
  W.times { |x| S[y][x] = (a += S[y][x]) }
end

W.times do |x, a = 0|
  H.times { |y| S[y][x] = (a += S[y][x]) }
end

puts H.times.collect { |y| W.times.collect { |x| S[y][x] } * " " }

参考

Discussion