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

AtCoder Beginner Contest 452

E - You WILL Like Sigma Problem

問題文

あなたの今日のラッキーギリシャ文字はシグマです。シグマを $2$ つも使ったこの問題を解けば、きっと幸運が舞い込むことでしょう。

長さ $N$ の正整数列 $A = (A_1, \cdots, A_N)$ および長さ $M$ の正整数列 $B = (B_1, \cdots, B_M)$ が与えられます。
$\displaystyle \sum_{i=1}^{N} \sum_{j=1}^{M} A_i \cdot B_j \cdot (i \bmod j)$ の値を $998244353$ で割ったあまりを求めてください。

制約

$1 \leq N,M \leq 5 \times 10^5$
$1 \leq A_i, B_j \leq 5 \times 10^5$
入力される値はすべて整数

解法

唐突に運勢を占われて戸惑う。

$\frac{N}{1}+\frac{N}{2}+\frac{N}{3}+...+\frac{N}{N}$ が $O(N \log{N})$ になるのって、 F～G問題くらいの知識だったイメージがあるけど、Eで出るんですねえ。

ひとまず、$(i,j)$ で表を作って、$(i \bmod j)$ がいくつになるかを表すと、

i\j  1  2  3  4  5  6  7  ... M
--------------------------------
1    0  1  1  1  1  1  1  ... 1
2    0  0  2  2  2  2  2  ... 2
3    0  1  0  3  3  3  3  ... 3
4    0  0  1  0  4  4  4  ... 4
5    0  1  2  1  0  5  5  ... 5
6    0  0  0  2  1  0  6  ... 6

$i=j$ の対角線を境に、右上側は全て $i$ のままとなる。
左下が $j$ 毎に様々な値を取るので厄介だが、これも $j$ 毎に $0,1,2,...,j-1$ を繰り返す周期性はある。

mod をとらない $\displaystyle \sum_{i=1}^{N} \sum_{j=1}^{M} A_i \cdot B_j \cdot i$ ならば、 $\displaystyle \left ( \sum_{i=1}^{N} A_i \cdot i \right ) \cdot \left ( \sum_{j=1}^{M} \cdot B_j \right )$ として、まとめて計算できる。そこからの差分を考えると、

Ai*Bj*i から減らさなければいけない Ai*Bj*x の x

 i\j  1  2  3  4  5  6  7  8  9 10
 ----------------------------------
 1    1
 2    2  2
 3    3  2  3
 4    4  4  3  4
 5    5  4  3  4  5
 6    6  6  6  4  5  6
 7    7  6  6  4  5  6  7
 8    8  8  6  8  5  6  7  8
 9    9  8  9  8  5  6  7  8  9
10   10 10  9  8 10  6  7  8  9 10

つまり、$j$ 毎に、$j$ の倍数の $i$ で $j$ だけ増えるような構造になっている。累積和を利用すれば、

 i\j   1    2    3    4    5    6   ...
 ---------------------------------------
 1   1*B1                               │
 2   1*B1 2*B2                          │各 j につき、
 3   1*B1      3*B3                     │i = jk（k=1,2,...）の位置に
 4   1*B1 2*B2      4*B4                │j*Bj を加算後、累積和
 5   1*B1                5*B5           │
 6   1*B1 2*B2 3*B3           6*B6      │
 7   1*B1                               │
 8   1*B1 2*B2      4*B4                │
 9   1*B1      3*B3                     │
10   1*B1 2*B2           5*B5           ↓

各 $i$ につき $A_i \times 累積和[i]$ だけ減じれば、答えとなる。

Python3

def solve(n, m, aaa, bbb):
    MOD = 998244353
    dec = [0] * (n + 2)
    for j in range(1, m + 1):
        for i in range(j, n + 1, j):
            dec[i] += bbb[j - 1] * j
    for i in range(1, n + 2):
        dec[i] = (dec[i - 1] + dec[i]) % MOD
    sb = sum(bbb)
    ans = 0
    for i in range(1, n + 1):
        ans += aaa[i - 1] * i * sb
        ans -= dec[i] * aaa[i - 1]
        ans %= MOD
    return ans


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

F - Interval Inversion Count

問題文

正整数 $N$ と、$(1,2,\ldots,N)$ の順列 $P=(P _ 1,P _ 2,\ldots,P _ N)$ が与えられます。
整数 $K$ が与えられます。次の $2$ つの条件を満たす整数の組 $(l,r)$ がいくつあるか求めてください。
- $1\le l\le r\le N$
- 列 $(P _ l,P _ {l+1},\ldots,P _ r)$ の転倒数が $K$ と等しい。

制約

$1\le N\le5\times10 ^ 5$
$0\le K\le\dfrac{N(N-1)}2$
$1\le P _ i\le N\ (1\le i\le N)$
$P _ i\ne P _ j\ (1\le i\lt j\le N)$
入力はすべて整数

解法

「転倒数が $x$ の区間を内包する区間」の転倒数は必ず $x$ 以上なので、 $[l,r)$ の $l$ を固定して $r$ を伸ばしていき、転倒数が $K+1$ 以上になったら打ち切ってよい。
さらにそこから $l$ を縮めていき、$K$ 以下になったら再び $r$ を伸ばし、、、という尺取法がベースとなる。

Fenwick Tree を使えば、ある数列の末尾に要素追加、または先頭から要素削除して、転倒数を更新する操作は、$O(\log{N})$ でおこなえる。

転倒数がちょうど $K$ になったところで答えに1ずつ加算していけば良さそうだが、、、

しかし、例えば以下のようなケースでは、

N=7  K=3
1 2 5 4 3 6 7

中央の $(5,4,3)$ の部分で転倒数 $K$ を満たしつつ、左は $1,2$、右は $6,7$ の任意の箇所まで伸ばしても転倒数は変化しないので、 $9$ 通りの区間が条件を満たすことになる。こういうのが、単純な尺取法では上手く処理できない。

尺取法で $r$ に相当するポインタを2つ持ち、

$r_1$: 転倒数がはじめて $K$ 以上となる位置
$r_2$: 転倒数がはじめて $K+1$ 以上となる位置

とすれば、各 $l$ を固定した場合につき、$r_2-r_1$ が答えとなる。

ただし、$K=0$ の時は $r_1$ が上手く働かないので、別に求める。

Python3

from atcoder.fenwicktree import FenwickTree


def solve(n, k, ppp):
    if k == 0:
        ans = 0
        bw = ''.join(str(int(ppp[i] < ppp[i + 1])) for i in range(n - 1))
        for chain in bw.split('0'):
            m = len(chain) + 1
            ans += m * (m + 1) // 2
        return ans

    fwt1 = FenwickTree(n)
    fwt2 = FenwickTree(n)
    ppp = [p - 1 for p in ppp]
    l, r1, r2 = 0, 0, 0
    inv1, inv2 = 0, 0
    ans = 0
    while l < n:
        while r1 < n and inv1 < k:
            p = ppp[r1]
            inv1 += fwt1.sum(p, n)
            fwt1.add(p, 1)
            r1 += 1

        while r2 < n and inv2 <= k:
            p = ppp[r2]
            inv2 += fwt2.sum(p, n)
            fwt2.add(p, 1)
            r2 += 1

        ans += r2 + int(inv2 == k) - r1

        p = ppp[l]
        fwt1.add(p, -1)
        fwt2.add(p, -1)
        inv1 -= fwt1.sum(0, p)
        inv2 -= fwt2.sum(0, p)
        l += 1

    return ans


n, k = map(int, input().split())
ppp = list(map(int, input().split()))
ans = solve(n, k, ppp)
print(ans)

G - 221 Substring

問題文

正整数からなる数列 $X = (X_1, \dots, X_n)$ に対し、その連長圧縮の任意の (長さ・値) ペアにおいて長さと値が等しいとき、かつそのときに限り、$X$ を 221 数列と呼びます。
- たとえば、$(2,2,3,3,3,1,2,2)$ は 221 数列ですが、$(1,1)$ や $(4,4,1,4,4)$ は 221 数列ではありません。
長さ $N$ の正整数列 $A = (A_1, \dots, A_N)$ が与えられます。$A$ の空でない連続部分列として現れる 221 数列としてあり得るものの個数を求めてください。
ただし、異なる位置から取り出した連続部分列であっても、数列として一致するものは区別せずまとめて $1$ 通りとして数えます。

制約

$1 \leq N \leq 500\,000$
$1 \leq A_i \leq 9$
入力される値はすべて整数

解法

「221文字列の種類数」の問題を、単なる「部分文字列の種類数」に言い換えるのはそこまで難しくない。
$A$ を連長圧縮して、(長さ $l$、値 $c$) の各ペアに付き、

5 5 5 のように、$l \lt c$ のとき、221文字列に含めることはできない。
3 3 3 のように、$l = c$ のとき、221文字列のどこに出現しても良い要素となる。
2 2 2 のように、$l \gt c$ のとき、221文字列の端にのみ出現できる要素となる。

よって、連長圧縮の $(l,c)$ を、$c$ に置換しつつ、$l \neq c$ の箇所で数列を区切っていく。

 5 5 5 3 3 3 2 2 2 1 1 1 4 4 4 4 1 2 2 8 2 2 2 2 3 3 3
                          ↓
      |  3   2 | 2 1 | 1    4    1  2  |  2 | 2    3

→ (3,2), (2,1), (1,4,1,2), (2), (2,3)

これらの複数の数列の中に登場する、ユニークな部分文字列の種類数を数えたい。

これはやや高度な典型問題で、Suffix Automaton が使える。

Suffix Automaton - よすぽの日記

または、公式Editorialのように、1～9以外の文字（たとえば0）で全て連結した上で Suffix Array + LCP からも求めることができる。

解説 - AtCoder Beginner Contest 452

Python3

from itertools import groupby


class SuffixAutomaton:
    def __init__(self):
        self.sa_len = [0]
        self.sa_link = [-1]
        self.sa_trans = [{}]
        self.size = 1

    def extend(self, last, c):
        # Generalized SAM: transition from last already exists
        if c in self.sa_trans[last]:
            q = self.sa_trans[last][c]
            if self.sa_len[q] == self.sa_len[last] + 1:
                return q
            clone = self.size
            self.sa_len.append(self.sa_len[last] + 1)
            self.sa_link.append(self.sa_link[q])
            self.sa_trans.append(self.sa_trans[q].copy())
            self.size += 1
            self.sa_link[q] = clone
            p = last
            while p != -1 and self.sa_trans[p].get(c) == q:
                self.sa_trans[p][c] = clone
                p = self.sa_link[p]
            return clone

        cur = self.size
        self.sa_len.append(self.sa_len[last] + 1)
        self.sa_link.append(0)
        self.sa_trans.append({})
        self.size += 1

        p = last
        while p != -1 and c not in self.sa_trans[p]:
            self.sa_trans[p][c] = cur
            p = self.sa_link[p]

        if p == -1:
            self.sa_link[cur] = 0
        else:
            q = self.sa_trans[p][c]
            if self.sa_len[q] == self.sa_len[p] + 1:
                self.sa_link[cur] = q
            else:
                clone = self.size
                self.sa_len.append(self.sa_len[p] + 1)
                self.sa_link.append(self.sa_link[q])
                self.sa_trans.append(dict(self.sa_trans[q]))
                self.size += 1
                self.sa_link[q] = clone
                self.sa_link[cur] = clone
                while p != -1 and self.sa_trans[p].get(c) == q:
                    self.sa_trans[p][c] = clone
                    p = self.sa_link[p]

        return cur

    def add_string(self, s):
        last = 0
        for c in s:
            last = self.extend(last, c)

    def count_distinct_substrings(self):
        return sum(
            self.sa_len[v] - self.sa_len[self.sa_link[v]]
            for v in range(1, self.size)
        )

    def debug_print(self):
        keys = sorted({k for t in self.sa_trans for k in t})
        header = f"{'i':>4} {'len':>4} {'link':>4}" + "".join(f" {k!s:>4}" for k in keys)
        print(header)
        print("-" * len(header))
        for i in range(self.size):
            trans = " ".join(
                f"{self.sa_trans[i].get(k, '-')!s:>4}" for k in keys
            )
            print(f"{i:>4} {self.sa_len[i]:>4} {self.sa_link[i]:>4} {trans}")


def solve(n, aaa):
    bbb = []
    sa = SuffixAutomaton()
    for c, itr in groupby(aaa):
        l = len(list(itr))
        if l < c:
            if len(bbb) > 0:
                sa.add_string(bbb)
            bbb.clear()
        elif l == c:
            bbb.append(c)
        else:
            bbb.append(c)
            sa.add_string(bbb)
            bbb = [c]
    if len(bbb) > 0:
        sa.add_string(bbb)

    ans = sa.count_distinct_substrings()
    return ans


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

目次

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

E - You WILL Like Sigma Problem

問題文

制約

解法

F - Interval Inversion Count

問題文

制約

解法

G - 221 Substring

問題文

制約

解法