🍡

[鉄則A08] Two Dimensional Sum

2024/04/25に公開

Ruby

二次元の累積和の基本問題。

問題を要約する

矩形範囲の合計値は？

解き方

累積和テーブルを作る
右下 + 左上 - 左下 - 右上

例1でシミュレーション

元のテーブルから、

	0	1	2	3	4
0	2	0	0	5	1
1	1	0	3	0	0
2	0	8	5	0	2
3	4	1	0	0	6
4	0	9	2	7	0

横方向の累積和を作って、

	1	2	3	4	5
0	0	0	0	0	0
1	2	2	2	7	8
2	1	1	4	4	4
3	0	8	13	13	15
4	4	5	5	5	11
5	0	9	11	18	18

それを縦に累積して、

	1	2	3	4	5
0	0	0	0	0	0
1	2	2	2	7	8
2	3	3	6	11	12
3	3	11	19	24	27
4	7	16	24	29	38
5	7	25	35	47	56

最後に矩形の4点の値

右下(角の内) → 38
左上(角の外) → 2
左下(辺の外) → 7
右上(辺の外) → 8

を見て、

右下 + 左上 - 左下 - 右上

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

38 + 2 - 7 - 8  # => 25

求まる。

単純な実装 (TLE)

H, W = gets.split.collect(&:to_i)                   # => [5, 5]
F = H.times.collect { gets.split.collect(&:to_i) }  # => [[2, 0, 0, 5, 1], [1, 0, 3, 0, 0], [0, 8, 5, 0, 2], [4, 1, 0, 0, 6], [0, 9, 2, 7, 0]]
Q = gets.to_i
Q = Q.times.collect { gets.split.collect(&:to_i) }  # => [[2, 2, 4, 5], [1, 1, 5, 5]]
Q = Q.collect { |e| e.collect(&:pred) }             # => [[1, 1, 3, 4], [0, 0, 4, 4]]
ans = Q.collect do |y1, x1, y2, x2|
  (x1..x2).sum do |x|
    (y1..y2).sum do |y|
      F[y][x]
    end
  end
end
ans                                                 # => [25, 56]
puts ans

わかりやすいが計算時間の問題で通らない。

解説の模範解答風 (AC)

H, W = gets.split.collect(&:to_i)                   # => [5, 5]
F = H.times.collect { gets.split.collect(&:to_i) }  # => [[2, 0, 0, 5, 1], [1, 0, 3, 0, 0], [0, 8, 5, 0, 2], [4, 1, 0, 0, 6], [0, 9, 2, 7, 0]]
Q = gets.to_i                                       # => 2
Q = Q.times.collect { gets.split.collect(&:to_i) }  # => [[2, 2, 4, 5], [1, 1, 5, 5]]

# ややこしくなるので 0-based に統一しておく
Q = Q.collect { |e| e.collect(&:pred) }             # => [[1, 1, 3, 4], [0, 0, 4, 4]]

# 上と左に一行一列入れるの重要
S = Array.new(H.next) { Array.new(W.next, 0) }      # => [[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 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 |y1, x1, y2, x2|
  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                                                 # => [25, 56]
puts ans

書き方で悩む

横方向の累積和を求める、

S = Array.new(H.next) { Array.new(W.next, 0) }
(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 = F.collect do |row|
  row.each_with_object([0]) { |e, m| m << m.last + e }
end
S.unshift(Array.new(W.next, 0))

と書けるが結局縦方向の累積和は、

(1..W).each do |x|
  (1..H).each do |y|
    S[y][x] += S[y.pred][x]
  end
end

と書かないといけないので x, y を添字にする書き方で統一した方がよさそう。

問題を要約する

解き方

例1でシミュレーション

単純な実装 (TLE)

解説の模範解答風 (AC)

書き方で悩む

Discussion