障害物の置かれたグリッド上で飛車が2回で移動できる範囲を求める。
解説
「→↓」で行けるマスの集合を $A$、「↓→」で行けるマスの集合を $B$ として、 $ | A \cup B | = | A | + | B | - | A \cap B | $ を計算する。
障害物に関する情報を以下の3つの配列で管理する。
- $\mathrm{row} [i] = $ 各行について障害物のある最も左の列番号
- $\mathrm{col} [j] = $ 各列について障害物のある最も上の行番号
- $\mathrm{obs} [i] = $ 各行の障害物のある列番号のリスト
$| A |$、$| B |$ はそれぞれ $| A | = \sum _ {i = 0} ^ {\mathrm{col} [0] - 1} \mathrm{row}[i]$、$| B | = \sum _ {j = 0} ^ {\mathrm{row} [0] - 1} \mathrm{col} [j]$ で計算できる。
$| A \cap B |$ は以下のように計算する。
- 列の長さ分の配列 $a$ を用意し、$0$ から $\mathrm{row}[0] - 1$ を $1$ で初期化する。
- 各行 $i = 0, \ldots, \mathrm{col[[0]}$ ごとに、 $\mathrm{obs} [i]$ に含まれる $j$ について $a[j]$ を $0$ に更新し、$\sum _ {j = 0} ^ {\mathrm{row}[i] - 1} a[j]$ を計算する。
配列の要素の更新と区間和を高速に計算するためには、Fenwick木を使う。
実装
Fenwick木
tree = FenwickTree(n) でサイズnの配列を0で初期化する。tree[i] で要素の取得、tree.sum(i, j) で [i, j) の区間和、tree.add(i, x) でi番目にxを加えることができる。
class FenwickTree: def __init__(self, n): self.size = n self.arr = [0] * (n + 1) def __getitem__(self, i): return self.sum(i + 1) - self.sum(i) def sum(self, i): s = 0 tmp = i while tmp: s += self.arr[tmp] tmp -= tmp & -tmp return s def add(self, i, x): tmp = i + 1 while tmp <= self.size: self.arr[tmp] += x tmp += tmp & -tmp
残り
h, w, m = map(int, input().split()) row, col = [w] * h, [h] * w obs = [[] for _ in range(h)] for _ in range(m): i, j = map(lambda x: int(x) - 1, input().split()) row[i] = min(row[i], j) col[j] = min(col[j], i) obs[i].append(j) ans = sum(row[: col[0]]) + sum(col[: row[0]]) tree = FenwickTree(w) for j in range(row[0]): tree.add(j, 1) for i in range(col[0]): for j in obs[i]: tree.add(j, -tree[j]) ans -= tree.sum(row[i]) print(ans)
Submission #18899800 - Panasonic Programming Contest (AtCoder Beginner Contest 186)