Branched Evolution

Competitive Programming in Python

ABC 186 F - Rook on Grid

障害物の置かれたグリッド上で飛車が2回で移動できる範囲を求める。

F - Rook on Grid

解説

「→↓」で行けるマスの集合を $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 |$ は以下のように計算する。

  1. 列の長さ分の配列 $a$ を用意し、$0$ から $\mathrm{row}[0] - 1$ を $1$ で初期化する。
  2. 各行 $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)