🍡

[鉄則B08] Counting Points

2024/04/26に公開

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

二次元の累積和の応用問題。

問題を要約する

  • 矩形で囲んだセルの合計値は?

解き方

  1. 座標のリストから二次元のテーブルを作る
  2. 累積和テーブルを作る
  3. 右下 + 左上 - 左下 - 右上

気をつける点

  • 躊躇せずすべてを 0-based に統一すべし

そのまま扱っていたら大混乱に陥った。

例1でシミュレーション

元の座標を、

[[1, 1], [3, 4], [4, 3]]

から、

[[0, 0], [2, 3], [3, 2]]

に補正してから次の二次元マップを作る。

0 1 2 3
0 1 0 0 0
1 0 0 0 0
2 0 0 0 1
3 0 0 1 0

すると基本問題と合流する。

あとは同じ流れで横方向の累積和を作り、

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

それを縦に累積して、

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

最後に矩形の4点の値

  • 右下(角の内) → 3
  • 左上(角の外) → 1
  • 左下(辺の外) → 1
  • 右上(辺の外) → 1

を見て、

  • 右下 + 左上 - 左下 - 右上

をすると不思議なことに矩形の合計値が

3 + 1 - 1 - 1  # => 2

求まる。

単純な実装 (TLE)

N = gets.to_i
XY = N.times.collect { gets.split.collect(&:to_i) }  # => [[1, 1], [3, 4], [4, 3]]
Q = gets.to_i
Q = Q.times.collect { gets.split.collect(&:to_i) }   # => [[2, 2, 4, 4]]
ans = Q.collect do |x1, y1, x2, y2|
  XY.count do |x, y|
    x.between?(x1, x2) && y.between?(y1, y2)
  end
end
ans                                                  # => [2]
puts ans

わかりやすいが TLE になる。

解説の模範解答風 (AC)

N = gets.to_i
XY = N.times.collect { gets.split.collect(&:to_i) }  # => [[1, 1], [3, 4], [4, 3]]
Q = gets.to_i                                        # => 1
Q = Q.times.collect { gets.split.collect(&:to_i) }   # => [[2, 2, 4, 4]]

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

# 問題毎にマップのサイズを決めたいから
W = XY.collect { |x, y| x }.max.next                 # => 4
H = XY.collect { |x, y| y }.max.next                 # => 4

# 座標のリストから二次元マップを作る
F = Array.new(H) { Array.new(W, 0) }
XY.each do |x, y|
  F[y][x] += 1
end
F                                                    # => [[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]

# 上と左に一行一列入れるの重要
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

# 矩形の和
ans = Q.collect do |x1, y1, x2, y2|
  rb = S[y2 + 1][x2 + 1]            # 右下
  lt = S[y1][x1]                    # 左上
  lb = S[y2 + 1][x1]                # 左下
  rt = S[y1][x2 + 1]                # 右上
  rb + lt - lb - rt                 # 右下 + 左上 - 左下 - 右上
end
ans                                                  # => [2]
puts ans

矩形の和を求める際に左上が角の外を参照する。つまり、はみ出ているため、上と左にパディングが必要になる。

Discussion