AtCoder Beginner Contest 291（Sponsored by TOYOTA SYSTEMS）G,Ex問題メモ

AtCoder Beginner Contest 291（Sponsored by TOYOTA SYSTEMS）

G - OR Sum

• 長さ $N$ の2つの整数列 $A=(A_1,...,A_N),B=(B_1,...,B_N)$ が与えられる
• 以下の操作を順に行う
  • $A$ を好きなだけ転回する
  • 転回後、$\displaystyle \sum_{i} (A_i \ | \ B_i)$ をスコアとする
    • ただし、$|$ はビット和
• スコアの最大値を求めよ
• $2 \le N \le 5 \times 10^5$
• $0 \le A_i,B_i \lt 2^5$
• 実行制限時間: 8sec.

いくつかの典型を組み合わせる。（よく同時に出てくる典型なので、まとめて1つの典型とも言えるかも）

$A_i,B_i$ の制約からも何となく察せられるが、bit毎に考えられる。

ただし、「1bit目が“1”になる要素数の最大値」「2bit目が“1”になる要素数の最大値」、、、を求めても、 その最大値を達成できる転回回数が異なっていては同時には達成できないため、転回回数は区別できる必要がある。

「$k$ 回転回したときの、$d$ bit目が“1”になる要素数」が求められれば、$k$ 毎にスコアを計算できる。

2つの数列 $a,b$ の畳み込み $c$ は、$\displaystyle c_k=\sum_{i+j=k}a_i b_j$、 つまり添字の和が $k$ になる全 $(i,j)$ の組についての $a_ib_j$ の合計となる。

ここで、片方（$b$ とする）をひっくり返すと、「添字の和が同じになる」から 「添字の差が同じになる」$(i,j)$ についての合計を求める問題に変換でき、 特に $|b|=M$ として $c_M$ が、$i=j$ の時の答えとなる。

本問題の場合は、転回するので「$A$ を2回繰り返した数列」と「$B$ をひっくり返した数列」の畳み込み結果の 添字 $N～2N-1$ の部分が、0回、1回、、、転回したときの畳み込み結果となる。

畳み込みは、各要素が $0$ か $1$ の場合は特に、「$a_i$ と $b_j$ がともに“1”である要素数」を表すともいえる。

本問題の場合は、逆にORを求めなければならないが、 0と1を反転させた上で畳み込みすることで、「$a_i$ と $b_j$ がともに“0”である要素数」を求めることができる。

よって、それを $N$ から引くことで、「$a_i$ または $b_j$ が“1”である要素数」を求めることができる。

Python3

import numpy as np


def convolve(f, g):
    fft_len = 1
    true_len = len(f) + len(g) - 1
    while fft_len < true_len:
        fft_len <<= 1

    Ff = np.fft.rfft(f, fft_len)
    Fg = np.fft.rfft(g, fft_len)
    Fh = Ff * Fg

    h = np.fft.irfft(Fh, fft_len)
    h = np.rint(h).astype(np.int64)

    return h[:true_len]


n = int(input())
aaa = list(map(int, input().split()))
bbb = list(map(int, input().split()))
aaa *= 2
bbb.reverse()

aaa = np.array(aaa, dtype=np.int64)
bbb = np.array(bbb, dtype=np.int64)
results = np.zeros(n, dtype=np.int64)

for b in range(5):
    fff = ((aaa >> b) & 1) ^ 1
    ggg = ((bbb >> b) & 1) ^ 1
    hhh = convolve(fff, ggg)
    results += (n - hhh[n:n * 2]) * (1 << b)

ans = results.max()
print(ans)

Ex - Balanced Tree

• $N$ 頂点の木 $T$ が与えられる
• 以下の条件をともに満たす $N$ 頂点の根付き木 $R$ を求めよ
  • 全ての頂点の組 $(x,y)$ に対し、次が成り立つ
    • $R$ 上での最小共通祖先 $LCA(x,y)=z$ のとき、$T$ 上で $x,y$ を結ぶパス上に $z$ が存在する
  • 根以外の全ての頂点 $v$ に対し、次が成り立つ
    • $R$ 上での $v$ を根とする部分木の頂点数を $size(v)$ で表す
    • この時、$size(v) \times 2 \le size(vの親)$ である
• $1 \le N \le 10^5$

問題の解釈が難しい。

木の再帰的な重心分解。
最近のARCで、重心を用いることが解法の1つとなる問題が出ていたので、記憶に新しい人が多かったのか、 解法のレア度の割に(?)正解者の多い結果となった。

まず、$T$ のままでも、適当な頂点を根として $R$ とすることで、1番目の条件は必ず満たされている。

2番目の条件と組み合わせたときにどうなるか、である。

2番目の条件は、

T
    1
   /|\
  2 8 9
  |    \
  3    10
 / \    |
4   5  11
|
6
|
7

みたいに、頂点1に対して「ある1つの子（上例では2）のサイズが、1自身+他の子（8や9）のサイズを全てあわせたのより大きいのはダメ」ということ。

かといって、「頂点7を8の下につなぎ替えてバランス取ります」みたいなことをしてしまうと、 1番目の条件が満たされなくなってしまう。 （つなぎ替えた結果、$LCA(6,7)=1$ となるが、$T$ 上では $6-7$ パス上に $1$ は無いためアウト）

特に、$T$ 上で最初から隣接している2頂点 $(x,y)$ は、「$T$ 上で $x,y$ を結ぶパス」が $x,y$ しかないため、 $R$ 上でも $LCA(x,y)=x または y$、つまり「どちらかがどちらかの祖先」である必要がある。

そのため、頂点を「ある部分木の子から、別の部分木の子へ」などと自由につなぎ替えることはできない。

$T$ の構造はなるべく維持したまま、上手くバランスを取る必要がある。

つなぎ替えるのがダメなら、回転させればいいじゃない。

一気に考えるのは難しいので、例えば $R$ の根を何にするべきかを考える。

「$T$ の構造を保ったまま、少なくとも、根と、その子にまつわる部分では条件2に違反しないように $R$ の根を決定する」には、$T$ の重心とするのがよい。

• ツリーの重心分解 (木の重心分解) の図解 - Qiita

重心は1つまたは2つ存在するが、2つの場合は適当にどちらか1つとすればよい。

上例では、1でなく2を根とすれば、少なくとも2とその子にまつわる部分では、条件2に違反することは無くなる。

     2
   ／ ＼
  3     1
 / \   / \
4   5 8   9
|         |
6        10
|         |
7        11

その各部分木も、再帰的に根が部分木の重心となるよう回転させてやれば、条件2が満たされる構造になることは分かる。

     2
   ／ ＼
  4     9
 / \   / \
6   3 1  10
|   | |   |
7   5 8  11

で、この時、条件1も実は上手くいくことが確認できる。
$T$ の重心が $c$ だったとして、

• $c$ の子で別の部分木に存在する2頂点 $x,y$ は、$T$ 上でも頂点 $c$ を経由しないといけないし、$R$ でも $LCA(x,y)=c$ が保たれている
• $c$ の子で同じ部分木に存在する2頂点は、再帰的に回転させた時、↑と同様のことが言える

1回の重心分解で各部分木は必ず半分以下になるため、1つの頂点が重心分解により探索されるのは $O(\log{N})$ 回。
重心分解は $O(N)$ なので、計算量は全体で $O(N \log{N})$ となる。

Python3

from collections import defaultdict


def get_centroid(n, links, s, removed):
    parents = {s: -1}
    subtree_sizes = {s: n}
    status = defaultdict(int)
    q = [s]
    limit = n // 2
    centroid = []
    while q:
        u = q[-1]
        if status[u] == 0:
            status[u] = 1
            for v in links[u]:
                if parents[u] == v:
                    continue
                if v in removed:
                    continue
                parents[v] = u
                q.append(v)
        else:
            status[u] = 2
            q.pop()
            my = 1
            is_centroid = True
            for v in links[u]:
                if status[v] != 2:
                    continue
                if v in removed:
                    continue
                scv = subtree_sizes[v]
                my += scv
                if scv > limit:
                    is_centroid = False
            subtree_sizes[u] = my
            if n - my > limit:
                is_centroid = False
            if is_centroid:
                centroid.append(u)
    return centroid, subtree_sizes


n = int(input())
links = [set() for _ in range(n)]
for _ in range(n - 1):
    u, v = map(int, input().split())
    u -= 1
    v -= 1
    links[u].add(v)
    links[v].add(u)

ans = [-1] * n
removed = set()

centroids, sizes = get_centroid(n, links, 0, removed)
root = centroids[0]
q = []
for v in links[root]:
    if sizes[v] < sizes[root]:
        q.append((root, v, sizes[v]))
    else:
        q.append((root, v, n - sizes[root]))
removed.add(root)

while q:
    p, v, m = q.pop()
    centroids, sizes = get_centroid(m, links, v, removed)
    c = centroids[0]
    ans[c] = p + 1
    for u in links[c]:
        if u in removed:
            continue
        if sizes[u] < sizes[c]:
            q.append((c, u, sizes[u]))
        else:
            q.append((c, u, m - sizes[c]))
    removed.add(c)

print(*ans)