AtCoder Beginner Contest 183 D,E,F問題メモ

AtCoder Beginner Contest 183

座標計算はゼロ除算をやらかしやすい。

D - Water Heater

問題

$N$ 人の人がいて、各々時刻 $[S_i, T_i)$ の間、毎時刻 $P_i$ リットルのお湯を使いたい
給湯器は毎時刻 $W$ リットルのお湯を供給できる
貯めておくことはできないとして、お湯の供給は間に合うか判定せよ
$1 \le N \le 2 \times 10^5$
$0 \le S_i \lt T_i \le 2 \times 10^5$

解法

区間更新・1点取得ができるデータ構造を用いる。

Segment Tree や Fenwick Tree 等、1点更新・区間取得ができるデータ構造を用いると、累積和からそれと同等のことが行える。今回はFenwick Treeで足りる。

全ての需要を反映（$[S_i,T_i)$ に $P_i$ を加算）した上で、各時刻 $t$ につき $W$ をオーバーしていないか確認していくとよい。

時刻の上限 $T_{max}=2 \times 10^5$ として $O((N+T_{max})\log{T_{max}})$ で行える。

Python3

class FenwickTree:
    def __init__(self, n):
        self.size = n
        self.tree = [0] * (n + 1)

    def sum(self, i):
        s = 0
        while i > 0:
            s += self.tree[i]
            i -= i & -i
        return s

    def add(self, i, x):
        while i <= self.size:
            self.tree[i] += x
            i += i & -i


n, w = map(int, input().split())
max_t = 200000
fwt = FenwickTree(max_t + 3)

for _ in range(n):
    s, t, p = map(int, input().split())
    fwt.add(s + 1, p)
    fwt.add(t + 1, -p)

for i in range(1, max_t + 2):
    if fwt.sum(i) > w:
        print('No')
        break
else:
    print('Yes')

そもそも今回は途中経過が必要ないので、imos法による累積和でよかったね。

解法2

「時刻 $S_i$ に $P_i$ 需要が増える」「$T_i$ に $P_i$ 減る」という計 $2N$ 個のイベントを時刻順にソートして、先頭から需要量をシミュレートしていってもよい。

注意点として、同時刻の減るイベントと増えるイベントでは、減るイベントを先に処理するようにする。（$T$ で終わる需要と $T$ で始まる需要は重複しないので）

この方法では時刻の長さ分の配列を用意する必要が無いので、時刻の取り得る範囲が $10^{18}$ など大きくても大丈夫。

Python3

n, w = map(int, input().split())
events = []

for _ in range(n):
    s, t, p = map(int, input().split())
    events.append((s, p))
    events.append((t, -p))

events.sort()
demand = 0
for _, p in events:
    demand += p
    if demand > w:
        print('No')
        break
else:
    print('Yes')

E - Queen on Grid

問題

$H \times W$ の盤面、'.'が通過可能マス、'#'が壁
盤面の左上 $(1,1)$ に、チェスのクイーンが置かれている
このクイーンは、壁にぶつかるまで、右、下、右下の直線方向へ1手で好きなだけ移動できる
$(H,W)$ への行き方が何通りあるか、$\mod{10^9+7}$ で求めよ
$2 \le H,W \le 2000$

解法

右方向に2マス移動するのにも、「1手で2マス移動する」「1マス、1マスと2手かける」の2通りがある点に注意。

DPをするのだが、工夫しないとTLEになる。

たとえば、以下のようなDPを作って

$DP[i][j]=$ マス $(i,j)$ で1手を終える場合の数

更新時、以下のようにあるマスから遷移できる全てのマスに個別に遷移していると、$O(HW(H+W))$ となり、間に合わない。

$DP[i+1][j]+=DP[i][j]$
$DP[i+2][j]+=DP[i][j]$
$DP[i+3][j]+=DP[i][j]$
…（壁にぶつかるまで）
$DP[i][j+1]+=DP[i][j]$
$DP[i][j+2]+=DP[i][j]$
$DP[i][j+3]+=DP[i][j]$
…（壁にぶつかるまで）
$DP[i+1][j+1]+=DP[i][j]$
$DP[i+2][j+2]+=DP[i][j]$
$DP[i+3][j+3]+=DP[i][j]$
…（壁にぶつかるまで）

同じ方向への移動を、上手くまとめてやる。

$DP[i][j][k]=$ マス $(i,j)$ を訪れる直前の移動が(k=0:右,1:下,2:右下)の場合の数
※そこで止まるとは限らない

とすると、3マスからのみの遷移で済む。もらうDPでの遷移例を書くと、

$DP[i][j][0] = DP[i][j-1][0] \times 2 + DP[i][j-1][1] + DP[i][j-1][2]$
$DP[i][j][1] = DP[i-1][j][1] \times 2 + DP[i-1][j][0] + DP[i-1][j][2]$
$DP[i][j][2] = DP[i-1][j-1][2] \times 2 + DP[i-1][j-1][0] + DP[i-1][j-1][1]$

つまり、例えば右移動で $(i,j)$ に至るのは $(i,j-1)$ からの遷移であり、その中でも「前と同じ方向（右）なら2倍」「異なる方向なら1倍」で足し合わせればよい。

───────┼──────────────┼─
(i, j-1)      │(i, j)                      │
  0.→: 5通り │  0.→: 5x2 + 2 + 8 = 20通り│
  1.↓: 2通り │  1.↓: ...                 │
  2.↘: 8通り │  2.↘: ...                 │

これは、以下を意味する。

前回と同じ方向の場合は、以下の2通り
- 直前のマスで止まる
- 止まらずに直接くる
対して異なる方向の場合は、必ず直前のマスで止まらなければならないので1通り

これで、$O(HW)$ となり、間に合う。

Python3

import os
import sys

import numpy as np


def solve(h, w, field):
    dp = np.zeros((h, w, 3), np.int64)
    MOD = 10 ** 9 + 7
    if field[0, 1] == 0:
        dp[0, 1, 0] = 1
    if field[1, 0] == 0:
        dp[1, 0, 1] = 1
    if field[1, 1] == 0:
        dp[1, 1, 2] = 1
    for i in range(h):
        for j in range(w):
            if field[i, j] == 1:
                continue
            if i > 0:
                dp[i, j, 1] += dp[i - 1, j, 1] * 2 + dp[i - 1, j, 0] + dp[i - 1, j, 2]
            if j > 0:
                dp[i, j, 0] += dp[i, j - 1, 0] * 2 + dp[i, j - 1, 1] + dp[i, j - 1, 2]
            if i > 0 and j > 0:
                dp[i, j, 2] += dp[i - 1, j - 1, 2] * 2 + dp[i - 1, j - 1, 0] + dp[i - 1, j - 1, 1]
            dp[i, j] %= MOD
    return dp[-1, -1].sum() % MOD


signature = '(i8,i8,i1[:,:])'
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve', signature)(solve)
    cc.compile()
    exit()

if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit(signature, cache=True)(solve)
    print('compiled', file=sys.stderr)

h, w = map(int, input().split())
field = np.array([['.#'.index(c) for c in input()] for _ in range(h)], np.int8)
ans = solve(h, w, field)
print(ans)

もしくは、累積和を管理してもよいみたい。まぁ実質的に同じこととなる。

F - Confluence

問題

$N$ 人の生徒、生徒 $i$ の所属クラスは $C_i$
はじめ、各生徒は1人ずつの $N$ 個の集合を形成している
$Q$ 個のクエリを処理せよ
クエリ1: '1 a b'
- 生徒 $a$ を含む集合と生徒 $b$ を含む集合をマージする
クエリ2: '2 x y'
- その時点で生徒 $x$ を含む集合に属するクラス $y$ の生徒数を出力する
$1 \le N,Q \le 2 \times 10^5$

解法

Union-Findの拡張。

Union-Find木

集合のリーダーに「自身の集合に各クラスが何人いるか」という情報をdict(連想配列)などで持たせておく。

最初は、1人1集合なので各 $i$ につき $\{C_i:1\}$ という情報を持っている。

Union-Find上でマージする際、このdictもマージする。リーダーの持つdict同士で一方にもう一方を合算する。

① {1:3, 2:5, 3:1}            ┬→  {1:4, 2:5, 3:6, 4:2, 5:1}
② {1:1,      3:5, 4:2, 5:1}  ┘

その際、dictの要素数が大きい方をベースとして、小さい方を統合する。

上例では、②をベースとし、①の方に対してイテレートを回して合算していく。

そうすることで、マージにかかる計算量は全体で $O(N\log{N})$ に抑えられる。
逆にこうしないと、意地悪なクエリでは最悪 $O(N^2)$ の計算量がかかる。

"データ構造をマージする一般的なテク" とは？ - (iwi) ｛反省します - TopCoder部

基準を要素数で無く集合サイズとした場合

上記では、どちらをマージのベースとするかについてdictの要素数を基準としたが、一般的なUnion-Findでも、集合のマージ時に「集合サイズを基準として大きい方に小さい方をマージする」というルールを使うことは多い。この基準をそのままdictのマージに流用しても、下記のように多少非効率が生じうるが、およそ同様の効率化が図られる。（敢えてそうする意味も薄いが、どちらをベースとするかが統一されるので混乱はしにくいかも）

① {1:1, 2:1, 3:1, ..., 9999:1}
② {1:10000}

これは、dictの要素数で考えると①の方をベースとすべきだが、
集合サイズを基準とすると②がベースとなり、多少非効率になる

しかし、①のようなデータを作るまでのマージや、
次にこのような非効率なマージが生じる場合を考えると、
非効率なマージが生じてしまう回数は限られると考えてよい

② {1:10001, 2:1, 3:1, ..., 9999:1}
③ {1:20000}
↑次に非効率なマージを生じさせるには、集合サイズが2倍以上になる必要がある

Python3

class UnionFind:
    def __init__(self, ccc):
        self.table = [{-1: -1, c: 1} for c in ccc]

    def _root(self, x):
        stack = []
        tbl = self.table
        while tbl[x][-1] >= 0:
            stack.append(x)
            x = tbl[x][-1]
        for y in stack:
            tbl[y][-1] = x
        return x

    def find(self, x, y):
        return self._root(x) == self._root(y)

    def unite(self, x, y):
        r1 = self._root(x)
        r2 = self._root(y)
        if r1 == r2:
            return
        d1 = self.table[r1]
        d2 = self.table[r2]
        if len(d1) >= len(d2):
            d1[-1] += d2[-1]
            d2[-1] = r1
            for cls, cnt in d2.items():
                if cls == -1:
                    continue
                if cls not in d1:
                    d1[cls] = cnt
                else:
                    d1[cls] += cnt
        else:
            d2[-1] += d1[-1]
            d1[-1] = r2
            for cls, cnt in d1.items():
                if cls == -1:
                    continue
                if cls not in d2:
                    d2[cls] = cnt
                else:
                    d2[cls] += cnt

    def get_class(self, x, y):
        r = self._root(x)
        d = self.table[r]
        if y in d:
            return d[y]
        else:
            return 0


n, q = map(int, input().split())
ccc = list(map(int, input().split()))
uft = UnionFind(ccc)
buf = []
for _ in range(q):
    op, a, b = map(int, input().split())
    if op == 1:
        uft.unite(a - 1, b - 1)
    else:
        buf.append(uft.get_class(a - 1, b))
print('\n'.join(map(str, buf)))

目次

AtCoder Beginner Contest 183 D,E,F問題メモ

D - Water Heater

問題

解法

解法2

E - Queen on Grid

問題

解法

F - Confluence

問題

解法

基準を要素数で無く集合サイズとした場合