CodeChef August Challenge 2018: Chef at the River - 日々ｄｒｄｒする人のメモ

Code Chef August Challenge 2018の問題: Chef at the River (Code: RIVER)
問題ページ: https://www.codechef.com/AUG18A/problems/RIVER/

問題概要

Chefと ${N}$ の動物が川の左岸におり、ボートを使って何回か往復することで全員右岸に渡りたい。
ボートは必ずChefが同行しなければならない。

${N}$ 匹の動物の中には、相性が悪いペアがおり、Chefと一緒にいない時に彼らが一緒にいてはいけない。
この悪いペアの関係は ${N}$ 頂点のグラフで表現され、 ${u}$ と ${v}$ が辺で繋がっている時、 ${u}$ , ${v}$ は相性が悪いことを表す。

最初は頂点1のみの根付き木になっており、 ${2, 3, ..., N}$ の頂点がこの根付き木に繋がっていく。
2の頂点が繋がった時から、 ${N}$ の頂点が繋がった時までの各 ${N-1}$ の状態について、制約を満たしながらChefと動物全員を対岸に渡すために必要な最小のボートの容量を求めよ。

制約

1つのテストケースに含まれるケース数: ${1 \le T \le 2 \times 10^5}$
${1 \le N \le 10^6}$
${1 \le p_i < i}$
全てのケースの ${N}$ の和は ${10^6}$ を超えない

部分点制約

(29点)
全てのケースの ${N}$ の和は ${1000}$ を超えない
(12点)
木の直径は40を超えない

解法

最終的な解法はHeavy-Light Decomposition + Binary Indexed Tree

探り

まずは何が解になるかを調べるために愚直解を実装して、根付き木の形と見比べてみる。
すると、解は根付き木の最小頂点被覆のサイズ(Chefを含むので+1 )と一致することが分かる。

部分点解法(29点, 12点)

この木に対する頂点被覆はdpで計算できる。
具体的には、頂点 ${v}$ 以下の部分木について、
${a_v}$ を頂点vを被覆しない時の最小頂点被覆数、 ${b_v}$ を頂点 ${v}$ を被覆する時の最小頂点被覆数とすると、
$a_v = \sum_{w} b_w$
$b_v = 1 + \sum_{w} \min(a_w, b_w)$
で計算できる。(wは頂点vの直近の子頂点)

これで各 ${N-1}$ の状態について木dpを計算すると ${O(N^2)}$ で29点がとれる。

さらに、 ${i-1}$ 番目の木と ${i}$ 番目の木で値が変化する箇所は、追加された頂点 ${v}$ から根頂点0までのパス上のみである。そのため、前の木のdpの値の状態から、追加された頂点 ${v}$ から根頂点0までの箇所をインクリメンタルに更新することで、下の図のように ${O(HN)}$ ( ${H}$ は木の高さ)で計算でき、さらに12点がとれる。
f:id:smijake3:20180814110456p:plain

満点解法(33点 + 26点)

ここから満点を取るには、この頂点 ${v}$ から根頂点0までの伝搬計算を高速に行う必要がある。
まず高速に更新する前に、計算しやすい形に変形する。

ここで、dpにおける各頂点の値を詳しく観察すると、頂点vと頂点vの親の関係にある頂点wについて、 ${\min(a_v, b_v) \le \min(a_w, b_w)}$ になる。このことから、 ${\min(a_v, b_v)}$ の値は根頂点0に到達し必ず解に加算されることが分かる。

ここで各頂点の値の状態を ${c_v = a_v - b_v}$ で表すことを考える。この場合、 ${\min(a_v, b_v) > 0}$ の値の情報は失われるが、上記の性質から、解の値を保持する変数ansに ${\min(a_v, b_v)}$ の値を加算することで対応する。

この ${c_v}$ の値を、追加された頂点vから根頂点まで伝搬させながら更新を行う。
この更新は以下の通りに行う。

新たに追加した頂点vの値 ${c_v}$ は-1にする
新たに追加された頂点vの親頂点の値 ${c_w}$ に+1を加算
頂点wの値に+1を加算する場合
- ${c_w}$ の元の値が-1であれば、親頂点の値に-1を加算し、ansの値を+1を加算する
- ${c_w}$ の元の値が-1でなければ、そこで更新を停止する
頂点wの値に-1を加算する場合
- ${c_w}$ の元の値が0であれば、親頂点の値に+1を加算し、ansの値を-1を加算する
- ${c_w}$ の元の値が0でなければ、そこで更新を停止する

例えば、以下の図のように、赤い頂点を追加する場合、各値は下図のように更新される。
f:id:smijake3:20180814110642p:plain
赤頂点の親頂点に+1されて、その親に-1、そのまた親に+1されて(値が-1ではないので)更新が停止する。
また、解変数ansには、 ${+1 - 1 = 0}$ が加算されて変化はない。

図や例から分かる通り、最初の親頂点は+1で更新を行い、-1, 0の交互になる限り+1, -1を交互に加算して更新、交互に-1, 0が出現しなくなる頂点で更新を停止する。

あとは、この更新を高速に行う。

まず、木において、各頂点vと根頂点0の間のクエリを扱うので、Heavy-Light Decompositionで木を分解し、heavy edgeで繋がった頂点で1まとまりにする。

そして、各頂点のまとまりについて2つのBITを使う。
1つ目は、頂点vの値が0か-1かを保持するBITである(具体的には各要素は0か1を持ち、累積和のMOD 2が0(0)か1(-1)で判別する)。このBIT上では、ある頂点から交互に0, -1が出現しなくなる頂点の特定と、特定の範囲の頂点の0, -1を-1, 0にflipする処理が ${O(\log N)}$ で実現できる。
2つ目は、頂点vの値が0より大きい場合にその値を持つBITである。このBIT上で二分探索することで、親頂点以降で0, -1以外の値が始めて出現する頂点の特定が ${O(\log N)}$ で実現できる。

各頂点を追加する際に、2つのBITから更新すべき頂点の範囲を特定し、その範囲の頂点の値の更新と、解の値ansの値の更新を行う。

この実装で満点が取れる。計算量は ${O(N \log^2N)}$

実装

部分点(29点)コード(PyPy2): Solution: 19610404 | CodeChef
部分点(41点)コード(PyPy2): Solution: 19611079 | CodeChef

提出コード(PyPy2): Solution: 19644615 | CodeChef

import sys
readline = sys.stdin.readline
write = sys.stdout.write

# Binary Indexed Tree

# 要素kまでの累積和S[k]を計算
def get(data, k):
    s = 0
    while k > 0:
        s += data[k]
        k -= k & -k
    return s
# k番目の要素v[k]の値をflipする (0 -> 1, 1 -> 0)
def flip(data, data0, k):
    l = len(data)
    if k >= l:
        return
    if data0[k]:
        # S[k]の値が1であれば1引く
        x = -1
    else:
        # S[k]の値が0であれば1足す
        x = 1
    data0[k] ^= 1
    if data0[k] == 2:
        x = -1
        data0[k] = 0
    while k < l:
        data[k] += x
        k += k & -k
# 通常のBITの加算 (S[k] += x)
def add(data, k, x):
    l = len(data)
    while k < l:
        data[k] += x
        k += k & -k

# w <= S[i] となる最小のインデックスiを求める
def lower_bound(data, w, b):
    l = len(data)
    if w <= 0:
        return -1
    x = 0
    while b:
        if x+b < l and data[x+b] < w:
            x += b
            w -= data[x]
        b >>= 1
    return x+1

# 要素v[i] ~ v[x]が全て1であるような最小のインデックスiを求める
# 累積和S[i] ~ S[x] (MOD 2)の値が0,1の交互になるような最小のインデックスiを求める
def lower_bound1(data, x, b):
    l = len(data)
    w = get(data, x)
    if w == x:
        return 0
    y = 0
    while b:
        if b < x and data[y+b] != b-x+w:
            y += b
            x -= b
            w -= data[y]
        b >>= 1
    return y+1
 
T = int(readline())
ans = []
for _ in xrange(T):
    N = int(readline())
    P = map(int, readline().split())
    G = [[] for i in xrange(N)]
    PR = [0]*N; PR[0] = -1
    for t in xrange(N-1):
        PR[t+1] = P[t]-1
        G[P[t]-1].append(t+1)
 
    # Heavy-Light Decomposition
    SZ = [0]*N
    NXT = [None]*N
    GR = []
    F = []; R = []
    for v in xrange(N-1, -1, -1):
        heavy = None; s = 0
        r = 1
        for w in G[v]:
            sz = SZ[w]
            if s < sz:
                heavy = w
                s = sz
            r += sz
        for w in G[v]:
            if w == heavy:
                NXT[v] = w
                continue
            F.append(w)
            R.append(v)
        SZ[v] = r
    F.append(0)
    R.append(-1)
 
    DS = []; DS0 = []; MS = []
    B2 = []
    LB = [0]*N; I = [0]*N
    for i, v in enumerate(F):
        LB[v] = i
        I[v] = cur = 1
        while NXT[v] is not None:
            v = NXT[v]
            cur += 1
            I[v] = cur
            LB[v] = i
        DS.append([0]*(cur+1))
        DS0.append([0]*(cur+1))
        MS.append([0]*(cur+1))
        B2.append(1 << cur.bit_length())
    B = [0]*len(DS)
    # Heavy-Light Decomposition ここまで
 
    # 頂点0の値(-1)を設定
    l = LB[0]; k = I[0]
    d = DS[l]; d0 = DS0[l]
    flip(d, d0, k)
    flip(d, d0, k+1)
 
    c = 0 # 今回の解の変数ans
    res = []
    for i in xrange(1, N):
        # 頂点e: まとめた頂点のうち、更新する際の基準となる頂点
        # v: 頂点eが持つべき値(0(0) or 1(-1))
        e = PR[i]; v = 1
 
        # heavy edgeでまとめた頂点単位で親に遷移しながら更新
        # cnt: この更新で値0, -1を-1, 0にflipした回数
        cnt = 0
        while e != -1:
            # このループ内は、1つのHL分解に含まれる頂点集合に対する処理
            # l: HL分解の中で処理する頂点集合のインデックス
            # k: 頂点集合内で、更新対象の頂点のうち最も深いところにある頂点の位置
            # d: 0, -1の値を持つBITデータ
            # m: 1以上の値を持つBITデータ
            l = LB[e]; k = I[e]
            d = DS[l]; m = MS[l]
            d0 = DS0[l]
            b2 = B2[l]
 
            # 値が0より大きくなる頂点が最初に出現する頂点位置xを求める
            x = lower_bound(m, get(m, k), b2)
 
            # -1, 0, -1, 0, ... が交互に出現しなくなる最初の頂点位置yを求める
            if (B[l] + get(d, k)) % 2 != v:
                # 頂点eが値vを持たない場合は位置yは頂点eになる
                y = k
            else:
                # 頂点eがvであればそれ以降の位置yを求める
                y = lower_bound1(d, k, b2)-1
                if y == d0[2] == 1:
                    y -= 1
 
            # 求めた2つの頂点のうち、最も近い方の頂点zまでが更新可能
            z = max(x, y)
            if z > 0:
                # 現在見ている頂点の中で、伝搬更新が停止する場合

                # flipする対象はz ~ kの間
                # 頂点zの値の偶奇も入れ替わるので、まとめてflipする
                flip(d, d0, z)
                flip(d, d0, k+1)

                cnt += k - z
 
                # 更新が停止する頂点zに対して更新する値を、1以上の値を持つBITにも加える
                if ((k - z) % 2)^v:
                    add(m, z, 1)
                else:
                    add(m, z, -1)
                break
            # 現在見ている頂点の中では、伝搬更新が停止しない場合
            # flipする対象は 1 ~ k
            B[l] ^= 1
            flip(d, d0, k+1)
            cnt += k
 
            # 次に更新する対象の頂点と、停止せず更新できる値を計算
            e = R[l]
            if k % 2 == 1:
                v ^= 1
 
        # 0, -1をflipした個数が奇数であれば解に+1する
        if cnt % 2:
            c += 1
 
        if i < 3:
            res.append("2")
        else:
            res.append(str(max(c+1, 3)))
 
        # 最後に追加する頂点の値(-1)を設定
        l = LB[i]; k = I[i]
        d = DS[l]; d0 = DS0[l]
        flip(d, d0, k)
        flip(d, d0, k+1)
 
    write(" ".join(res))
    write("\n")