AtCoder Beginner Contest 438 E,F,G問題メモ

AtCoder Beginner Contest 438

後半は、しっかり実装が求められる問題が多かった印象。

E - Heavy Buckets

• $N$ 人の人がおり、$N$ 個のバケツがあります。
• 人 $i$ ははじめバケツ $i$ のみを持っており、バケツ $i$ には何も入っていません。
• これから、以下の操作を $10^9$ 回行います。
  • $i = 1, 2, \ldots, N$ について同時に、人 $i$ が持っているバケツすべてに水を $i$ ずつ入れ、それらのバケツを人 $A_i$ に渡す。
• ただし、バケツに入れることのできる水の量に制限はありません。

• $i = 1, 2, \ldots, Q$ について以下の形式のクエリに答えてください。
  • $T_i$ 回目の操作の直後にバケツ $B_i$ に入っている水の量を求めてください。

• $2 \leq N \leq 2 \times 10^5$
• $1 \leq Q \leq 2 \times 10^5$
• $1 \leq A_i \leq N$
• $1 \leq T_i \leq 10^9$
• $1 \leq B_i \leq N$
• 入力される値はすべて整数

どの頂点からも $1$ 本ずつ辺が出ているようなグラフを Functional Graph と呼ぶ。

• 参考: AtCoder ABC 357 E - Reachability in Functional Graph (1D, 水色, 450 点) - けんちょんの競プロ精進記録

Functional Graph で○○回移動後の状態、というのはダブリングで管理できることが多い。 例えば今回は以下が前計算できると良さそうである。

• $\mathrm{Owner}[i,t]:=$ バケツ $i$ の $2^t$ 回の操作後の所有者
• $\mathrm{Water}[i,t]:=$ バケツ $i$ に $2^t$ 回の操作後に入っている水

ただ、Owner（に相当する情報）をダブリングで管理する問題はよく見るが、Waterは見慣れない。 一応、本当にダブリングで管理できるかどうか確認する。

• 操作回数によって、入る水の量は変わらない。つまり、任意の $x$ について、以下は変わらない
  • 最初に人 $y$ が持っているバケツに、$2^d$ 回の操作後に入っている水
  • $x$ 回の操作後に人 $y$ が持っているバケツに、$x+2^d$ 回の操作後に追加で入っている水
• 足し算なので、結合法則が成り立つ
  • あるバケツに $t$ 回目の操作で入る水の量を $W_t$ として、たとえば
  • $W_1+W_2+W_3+W_4=(W_1+W_2)+(W_3+W_4)$
  • のように、部分的な集計値同士を足し合わせても結果は変わらない

よって、Waterも問題なく管理できる。前計算は以下のようにおこなえる。計算量は $O(N \log{T})$ となる。

• $\mathrm{Owner}[i,t+1]=\mathrm{Owner}[\mathrm{Owner}[i,t],t]$
• $\mathrm{Water}[i,t+1]=\mathrm{Water}[i,t]+\mathrm{Water}[\mathrm{Owner}[i,t],t]$

クエリでは、たとえば $T_i=25=2^0+2^3+2^4$ のように2の冪乗に分解できるとき、 $2^0,2^3,2^4$ 回の操作をそれぞれ一気に進める（$t=0,3,4$ についてテーブルを参照する）ことで、1クエリ $O(\log{T})$ で答えを求められる。

Python3

def solve(n, q, aaa, queries):
    doubled_water = [list(range(1, n + 1))]
    doubled_owner = [[a - 1 for a in aaa]]
    for _ in range(30):
        no = []
        nw = []
        cw = doubled_water[-1]
        co = doubled_owner[-1]
        for i in range(n):
            no.append(co[co[i]])
            nw.append(cw[i] + cw[co[i]])
        doubled_water.append(nw)
        doubled_owner.append(no)
    buf = []
    for t, b in queries:
        b -= 1
        ans = 0
        i = 0
        while t:
            if t & 1:
                ans += doubled_water[i][b]
                b = doubled_owner[i][b]
            t >>= 1
            i += 1
        buf.append(ans)
    return buf


n, q = list(map(int, input().split()))
aaa = list(map(int, input().split()))
queries = [list(map(int, input().split())) for _ in range(q)]
ans = solve(n, q, aaa, queries)
print('\n'.join(map(str, ans)))

F - Sum of Mex

• $N$ 頂点の木 $T$ が与えられます。
• 各頂点には $0$ から $N-1$ までの番号がついており、$i$ 番目 $(1\le i\le N-1)$ の辺は頂点 $u_i$ と頂点 $v_i$ を双方向に繋いでいます。（頂点番号は 0-indexed で辺番号が 1-indexed であることに注意してください。）
• $0$ 以上 $N$ 未満の整数の組 $(i,j)$ に対し $f(i,j)$ を以下のように定めます：
  • 木 $T$ の頂点 $i$ から頂点 $j$ までのパスの頂点番号に 含まれない 最小の非負整数
• 木 $T$ の頂点 $i$ から頂点 $j$ までのパスに頂点 $i,j$ も含まれることに注意してください。

• $\displaystyle \sum_{0\le i \le j \lt N} f(i,j)$ の値を求めてください。

• $2\le N\le 2\times 10^5$
• $0\le u_i \lt v_i \lt N$
• 入力で与えられるグラフは木
• 入力される値は全て整数

この数え上げは、主客転倒っぽく、以下の合計として考えられる。

• 頂点 ⓪ を含むパスの個数
• 頂点 ⓪① を全て含むパスの個数
• 頂点 ⓪①② を全て含むパスの個数
• ：

⓪を含まないパスはMEXも $0$ なので無視できる。含むパスは最低、$1$ は保証される。
そのうち、①も含んでいたら最低 $2$ は保証されるので、さらにパスの個数だけ $1$ ずつ加算できる。
という考察を繰り返していくと、求める値が上と等しくなることが分かる。

⓪を根として、木DPで以下の情報を求めておく。

• $C_v:=$ 頂点 $v$ 以下の部分木サイズ
• $P_v:=$ 頂点 $v$ の親

⓪を端点に持つか、⓪の異なる子以下から2個選ぶ個数として求められる。$C_v$ があれば求められる。

①から親を辿っていく。いずれ⓪に行き着く。行き着く直前の⓪の子を $v$ とする。

     ⓪
   ／｜＼
 (v) ○ ○
  :  ： ：
 ↑
 ①
 /\
：：

以下の2つから1頂点ずつ選んで結んだパスが、⓪①を含むパスに該当する。

• ⓪および、⓪の(v)以外の子以下の頂点の個数（$C_0-C_v$）
• ①以下の頂点の個数（$C_1$）

2つの積がこの場合の答えとなる。

「⓪①$...k$ を含む最短パス」とは、⓪①$...k$ のうちいずれか2つが端点であり、残りがその上に乗っているようなパスを指すとする。
頂点の配置によっては存在しない場合もある。$k$ で存在しなかったら、$k+1$ 以降も当然、存在しない。

$D_k$ を「⓪①$...k$ を含む最短パスの個数」とする。$k=2,3,...$ について、存在する限り、$D_k$ を答えに加算していきたい。

最短パスの端点2つを $a,b$ とすると、（存在すれば）$D_k=C_a \times C_b$ で求められる。
（ただし端点が⓪の時のみ、先ほどの通り $C_v$ を除く必要がある。以降、$C_0$ としたときは除去後の値とする）

最初、$a=0,b=1$ とする。また、最短パス上の頂点を $S$ として記録しておく。$S$ の初期値は先ほど①から親を辿った過程で記録しておく。
$a,b,S$ を管理しつつ $k=2,3,...$ を順次パスに加えた結果を更新していく。

例えば $k=2$ の時を考える。

② が $S$ に含まれるなら、最短パスの端点 $a,b$ は変わらない。$a=0,b=1$ のままである。

それ以外の場合、②からも親を辿っていく。最初に以下のいずれかにぶつかったら止める。

• ⓪
  • $a=2$ に更新される。更新した上で $C_aC_b$ を答えに加える。
• ①
  • $b=2$ に更新される。更新した上で $C_aC_b$ を答えに加える。
• それ以外の $S$ に含まれる頂点
  • ⓪①②を全て含むパスは作れないことが確定する。
  • 当然、それ以降の③,④,…も作れないので、打ち切ってその時点の答えを出力する。

また、⓪か①に到達した場合は、そこまでに通過した頂点も $S$ に加えておく。

③以降も、この繰り返しでよい。つまり、

• $S$ に含まれていたら直前と同じ
• 親を辿っていき、$a$ か $b$ に辿り着いたら、辿り着いた方を更新
• それ以外の $S$ に辿り着いたら打ち切り

としていくと解ける。

最初の木DPでも、後の最短パスを更新していく過程でも、1頂点は高々定数回しか参照されないため、計算量は $O(N)$ となる。

Python3

def calc_subtree_sizes(n, links):
    sizes = [-1] * n
    parents = [-1] * n
    q = [0]
    while q:
        u = q[-1]
        if sizes[u] == -1:
            sizes[u] = -2
            for v in links[u]:
                if sizes[v] == -1:
                    q.append(v)
                    parents[v] = u
            continue
        q.pop()
        sizes[u] = sum(sizes[v] for v in links[u] if sizes[v] > 0) + 1
    return sizes, parents


def solve(n, edges):
    links = [[] for _ in range(n)]
    for u, v in edges:
        links[u].append(v)
        links[v].append(u)

    sizes, parents = calc_subtree_sizes(n, links)
    # print(parents)

    ans = n  # 0が端点

    # 0を通過
    tmp = 0
    for v in links[0]:
        ans += tmp * sizes[v]
        tmp += sizes[v]
    # print(ans)

    # 0-1
    on_path = {0}
    v = 1
    last_v = 1
    while v != 0:
        on_path.add(v)
        last_v = v
        v = parents[v]

    side_u = n - sizes[last_v]
    side_v = sizes[1]
    cur_u = 0
    cur_v = 1
    ans += side_u * side_v

    for i in range(2, n):
        if i in on_path:
            ans += side_u * side_v
            continue
        j = i
        tmp_path = set()
        make_branch = False
        while True:
            tmp_path.add(j)
            if j == cur_u:
                side_u = sizes[i]
                cur_u = i
                ans += side_u * side_v
                break
            if j == cur_v:
                side_v = sizes[i]
                cur_v = i
                ans += side_u * side_v
                break
            if j in on_path:
                make_branch = True
                break
            j = parents[j]
        if make_branch:
            break
        on_path.update(tmp_path)

    return ans


n = int(input())
edges = [list(map(int, input().split())) for _ in range(n - 1)]
ans = solve(n, edges)
print(ans)

G - Sum of Min

• 整数 $N,M,K$ と長さ $N$ の整数列 $A=(A_0,A_1,\ldots,A_{N-1})$ 、長さ $M$ の整数列 $B=(B_0,B_1,\ldots,B_{M-1})$ が与えられます。添字が $0$ から始まることに注意してください。
• $\displaystyle\sum_{i=0}^{K-1} \min(A_{i\bmod N}, B_{i \bmod M}) $ を $998244353$ で割ったあまりを求めてください。

• $1\le N,M\le 2\times 10^5$
• $1\le K\le 10^{18}$
• $1\le A_i,B_i\le 10^9$
• 入力される値は全て整数

ループするindexを $A$ も $B$ も1ずつ進めていくという問題。

とりあえず、後の計算が考えやすくなるので、$\mathrm{GCD}(N,M)=1$ とする。

そうで無い場合は、相対するペアが gcd ごとに飛び飛びになることを利用すれば、{gcd}個の問題に分割できる。
この時、$K$ はそれぞれ $\lfloor \frac{K}{gcd} \rfloor$ に加え、端数は左から順に1ずつ分配される。

N=6  M=9  K=50
A  1  2  3  4  5  6             GCD(6,9)=3
B  1  2  3  4  5  6  7  8  9

Aの(1,4)とペアになりうるのはBの(1,4,7)
Aの(2,5)とペアになりうるのはBの(2,5,8)
Aの(3,6)とペアになりうるのはBの(3,6,9)
↓
以下の3つの問題に分割できる。
K/gcd = 16 あまり 2 なので、分割後の K は、左から2個は17、残りは16となる。

N=2 M=3 K=17    N=2 M=3 K=17    N=2 M=3 K=16
A  1  4         A  2  5         A  3  6
B  1  4  7      B  2  5  8      B  3  6  9

以降、$\mathrm{GCD}(N,M)=1$ とする。

各 $i$（$A$ 側のindex）ごとに、ペアとなる $j$（$B$ 側のindex）は、以下のようになる。

N=7  M=5  K=60

A  0  1  2  3  4  5  6
B  0  1  2  3  4

   →順にペアとなるj
i↓                                     ペアになる順
 0:  0  2  4  1  3  0  2  4  1 | 3 ...    1  8 15
 1:  1  3  0  2  4  1  3  0  2 | 4 ...    2  9 ：
 2:  2  4  1  3  0  2  4  1  3 | 0 ...    3 10 ：
 3:  3  0  2  4  1  3  0  2  4 | 1 ...    4 11
 4:  4  1  3  0  2  4  1  3 ,--' 2 ...    5 12
 5:  0  2  4  1  3  0  2  4 |1   3 ...    6 13
 6:  1  3  0  2  4  1  3  0 |2   4 ...    7 14
                         K=60はここまで

以下の3つのフェイズ別に答えを求める。（用語は適当。後で説明）

• 完全ループ
• 部分ループ
• 端数

$NM$ 回の操作後、ペアが最初に戻り、ループする。これを“1完全ループ”とする。
完全ループする過程で、$\{A_0,...,A_{N-1}\}$ は、それぞれ $\{B_0,...,B_{M-1}\}$ と1回ずつペアになる。

   →順にペアとなるj
i↓  _____________
 0: |0  2  4  1  3| 0  2  4  1 | 3 ...
 1: |1  3  0  2  4| 1  3  0  2 | 4 ...
 2: |2  4  1  3  0| 2  4  1  3 | 0 ...
 3: |3  0  2  4  1| 3  0  2  4 | 1 ...
 4: |4  1  3  0  2| 4  1  3 ,--' 2 ...
 5: |0  2  4  1  3| 0  2  4 |1   3 ...
 6: |1  3  0  2  4| 1  3  0 |2   4 ...
     ~~~~~~~~~~~~~
      1完全ループ

よって、1完全ループの寄与は、ソートした $B$ 上で各 $A_i$ を二分探索して、 「$A_i$ 以下の $B_j$ の総和」+「$A_i \times$（$A_i$ より大きい $B_j$ の個数）」として求められる。

ループ回数だけ倍加して、答えに加算する。

以降、$K←(K \mod {NM})$ に更新する。

$B$ での位置はずれていてもよいが、$A$ が1周するまでを“1部分ループ”とする。ここが一番テクニカル？

   →順にペアとなるj
i↓                 _______
 0:  0  2  4  1  3 |0  2  4| 1 | 3 ...
 1:  1  3  0  2  4 |1  3  0| 2 | 4 ...
 2:  2  4  1  3  0 |2  4  1| 3 | 0 ...
 3:  3  0  2  4  1 |3  0  2| 4 | 1 ...
 4:  4  1  3  0  2 |4  1  3|,--' 2 ...
 5:  0  2  4  1  3 |0  2  4||1   3 ...
 6:  1  3  0  2  4 |1  3  0||2   4 ...
                    ~~~~~~~
                   3部分ループ

部分ループ回数を、$d = \lfloor \frac{K}{N} \rfloor$ とする。上例では $d=3$ である。

まず、$A_0$ は $d=3$ 回の部分ループ中で $B_0,B_2,B_4$ とペアになる。
この中だけで、完全ループの時のように「$A_i$ 以下の $B_j$ の総和」+「$A_i \times$（$A_i$ より大きい $B_j$ の個数）」を求めたい。

$A$ と $B$ に出てくる値で座標圧縮した上で、 有効な $B_i$ の「個数」と「総和」をFenwickTreeなどで管理すれば、求められる。

次 $A_1$ を求めようとすると、$B_1,B_3,B_0$ とペアになるので、$A_0$ の時からFenwickTreeの多くの要素について入れ替わりが発生してしまい、更新コストが大きくなる。
一方、$A_2$ なら、$B_2,B_4$ は保った上で、$B_0$ を除いて $B_1$ を追加するだけで更新できる。

このように、$t=0,1,...,M-1$ について、$i=tN \mod{M}$ とした $i$ の答えを求めていくようにすると、 FenwickTreeの更新箇所は1回につき $O(1)$ で抑えられる。

$N \gt M$ の場合、$i=0$ のタイミングで、$i=5$ など $i+nM$ で表される $i$ についても同時に求められる。

これで、全体 $O(N \log{N})$ などで総和が求められる。

あとは愚直に計算しても $O(N)$ である。

   →順にペアとなるj
i↓                          _
 0:  0  2  4  1  3  0  2  4 |1|| 3 ...
 1:  1  3  0  2  4  1  3  0 |2|| 4 ...
 2:  2  4  1  3  0  2  4  1 |3|| 0 ...
 3:  3  0  2  4  1  3  0  2 |4|| 1 ...
 4:  4  1  3  0  2  4  1  3 ,--' 2 ...
 5:  0  2  4  1  3  0  2  4 |1   3 ...
 6:  1  3  0  2  4  1  3  0 |2   4 ...

Python3

from bisect import bisect
from itertools import accumulate
from math import gcd

from atcoder.fenwicktree import FenwickTree


def solve(n, m, k, aaa, bbb):
    MOD = 998244353

    if n < m:
        n, m = m, n
        aaa, bbb = bbb, aaa

    def sub(n, m, k, aaa, bbb):
        # n, m 互いに素の場合
        lcm = n * m
        loop, k = divmod(k, lcm)
        ans = 0

        if loop > 0:
            ccc = sorted(bbb)
            acc = [0] + list(accumulate(ccc))
            for a in aaa:
                i = bisect(ccc, a)
                ans += a * (m - i)
                ans += acc[i]
                ans %= MOD
            ans *= loop
            ans %= MOD

        d, k = divmod(k, n)

        if d > 0:
            values = sorted(set(aaa) | set(bbb))
            mapping = {a: i for i, a in enumerate(values)}
            size = len(values)

            fwt_cnt = FenwickTree(size)
            fwt_sum = FenwickTree(size)
            for t in range(d):
                j = (t * n) % m
                b = bbb[j]
                bi = mapping[b]
                fwt_cnt.add(bi, 1)
                fwt_sum.add(bi, b)
            for t in range(m):
                s = (t * n) % m
                for i in range(s, n, m):
                    a = aaa[i]
                    ai = mapping[a]
                    cnt = fwt_cnt.sum(0, ai)
                    total = fwt_sum.sum(0, ai)
                    tmp = a * (d - cnt) + total
                    ans += tmp
                    ans %= MOD
                b = bbb[s]
                bi = mapping[b]
                fwt_cnt.add(bi, -1)
                fwt_sum.add(bi, -b)
                j = ((t + d) * n) % m
                b = bbb[j]
                bi = mapping[b]
                fwt_cnt.add(bi, 1)
                fwt_sum.add(bi, b)

        for i in range(k):
            j = (d * n + i) % m
            ans += min(aaa[i], bbb[j])
        ans %= MOD

        return ans

    g = gcd(n, m)
    n2 = n // g
    m2 = m // g
    k2, rem = divmod(k, g)
    ans = 0
    for i in range(g):
        ccc = [aaa[j] for j in range(i, n, g)]
        ddd = [bbb[j] for j in range(i, m, g)]
        ans += sub(n2, m2, k2 + int(i < rem), ccc, ddd)
        ans %= MOD

    return ans


n, m, k = list(map(int, input().split()))
aaa = list(map(int, input().split()))
bbb = list(map(int, input().split()))
ans = solve(n, m, k, aaa, bbb)
print(ans)