AtCoder Beginner Contest 444 E,F問題メモ

AtCoder Beginner Contest 444

0120-444-444は再春館製薬所の注文受付窓口。

E - Sparse Range

問題文

長さ $N$ の整数列 $(A_1,\dots,A_N)$ と正整数 $D$ が与えられます。
以下の条件をともに満たす整数の組 $(L,R)$ の個数を求めてください。
- $1 \leq L \leq R \leq N$
- $(A_L,A_{L+1},\dots,A_R)$ のどの $2$ つの要素も差が $D$ 以上である
  - すなわち、 $L \leq i \lt j \leq R$ を満たす全ての整数の組 $(i,j)$ について、 $|A_i-A_j|\geq D$ である

制約

$2\leq N \leq 4\times 10^5$
$1 \leq A_i \leq 10^9$
$1 \leq D \leq 10^9$

解法

尺取法。

$l=1,r=1$ とする。また、今の値の集合を $S=\{A_l,...,A_{r-1}\}$ とする。
以下を $l \le N$ まで繰り返す。
- $r$ を順に1つずつ進め、$A_r$ を $S$ に加えてもよいか調べる。条件を満たさなくなったら $S$ に加えず止める。
- 答えに $r-l$ を加算する（※）
- $l$ を1つ進め、$A_l$ を $S$ から除く。

（※）$L=l$ と固定した時、$R=l,l+1,...,r-1$ は全て条件を満たし、 $R \ge r$ 以降は満たさないので、$r-l$ 個だけ答えに寄与する。

「条件を満たさなくなった」、つまり「現在の $S$ の中に、$A_r$ との差が $D$ 未満のものが存在する」のをどう判定するか。

$S$ を SortedSet などで管理すれば、$A_r$ の次に小さい/大きい数を取得し差が $D$ 未満かどうかを判定できる。

また、SortedSet が重たければ、座標圧縮して FenwickTree などでも判定できる。

各 $i$ につき、$A_i-D,A_i,A_i+D$ を全て集めた集合（$\bigcup_{i=1}^{N} \{ A_i - D,\ A_i,\ A_i + D \}$）で座標圧縮しておき、FenwickTree 等に載せる。
$A_i$ を $S$ に加えるときは、$A_i$ の位置に1加算する。
$A_i$ を加えても良いか判定するときは、$(A_i-D,A_i+D)$（両端含まず）の和が0か調べる。

Python3

from atcoder.fenwicktree import FenwickTree


def solve(n, d, aaa):
    values = set()
    for a in aaa:
        values.add(a - d)
        values.add(a)
        values.add(a + d)
    values = sorted(values)
    mapping = {a: i for i, a in enumerate(values)}

    fwt = FenwickTree(len(values))
    l = 0
    r = 0
    ans = 0
    while l < n:
        while r < n:
            a = aaa[r]
            li = mapping[a - d]
            ri = mapping[a + d]
            if fwt.sum(li + 1, ri) > 0:
                break
            fwt.add(mapping[a], 1)
            r += 1

        ans += r - l

        la = aaa[l]
        fwt.add(mapping[la], -1)
        l += 1

    return ans


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

F - Half and Median

問題文

$N$ 本の棒があり、$i$ 本目の棒の長さは $A_i$ です。
以下の操作を $M$ 回行います。（必ず行えるような入力が与えられます）
- 長さが $2$ 以上の棒を $1$ 本選び、この棒の長さを $l$ と置く。この棒を、長さが $\left\lfloor\frac{l}{2}\right\rfloor$ と $\left\lceil\frac{l}{2}\right\rceil$ の $2$ 本の棒に分割する。
$M$ 回の操作後、$N+M$ 本の棒が残りますが、これらの棒の長さの中央値としてありうる値の最大値を求めてください。

$1$ つの入力につき、$T$ 個のテストケースを解いてください。

制約

$1 \leq T \leq 10^5$
$1 \leq N \leq 10^5$
$1 \leq A_i \leq 10^9$
$1 \leq M \leq \sum_{i=1}^{N}{A_i} - N$
$N+M$ は奇数
全てのテストケースにおける $N$ の総和は $10^5$ 以下
入力される値は全て整数

解法

非常に場合分けが面倒くさい。

中央値を求める問題は、二分探索で解けることが多い。

「中央値を $k$ 以上にすることが可能か？」の判定問題を考えることとする。
これは「$M$ 回の操作後、長さ $k$ 以上の棒を $\dfrac{N+M+1}{2}$ 本以上残すことが可能か？」と言い換えられる。

以下、長さ $k$ 以上の棒の本数の目標値を $X = \dfrac{N+M+1}{2}$ とする。

判定方法

（高速化はひとまず気にせず愚直でも良いので）判定方法を考える。

操作によって生じることは、以下の通りである。

長さが $\left\lfloor\frac{l}{2}\right\rfloor \ge k$ の棒を選び分割すると、長さ $k$ 以上の本数が1増える
長さが $\left\lfloor\frac{l}{2}\right\rfloor \lt k$ の棒を選び分割すると、長さ $k$ 以上の本数は増えないが、1回操作を消費できる

特に、たとえば $k=10$ で $l=19$ など「$k-1$ と $k$」に分かれるときは、長さ $k$ 以上の本数が増えない側である点に留意する。

まず、明らかに不可能な場合を除く。初期状態で $A_i \ge k$ である本数を $x$ として、

①1回の操作で長さ $k$ 以上の棒は高々1本しか増えないので、$x+M \lt X$ なら不可能。

それ以外の場合、実際に $A_i \ge k$ である各棒を「次に分割したら $k$ を下回ってしまう」まで分割する。

②この時点で長さ $k$ 以上の本数が $X$ 未満ならやっぱり不可能。

そうで無い場合、$x+M \ge X$ より、長さ $k$ 以上の棒を $X$ 本作ることはできる。だが、

③$M$ が大きすぎて「長さ $k$ 以上の棒 $X$ 本以外を全て長さ $1$ にしても、まだ操作が必要」な場合は不可能
④それより前に $M$ を消費し切れた場合に、やっと可能

となる。

②の時点で長さ $k$ 以上の本数が $X$ より多い場合は、 ③の判定で有利となるよう、$X$ 本残してあとは長さ1に分割してしまってよい。
その場合、長さ $l$ の棒は $l-1$ 回操作できるので、回数を多く消費するためには長い方から分割するのがよい。

(例) 長さ $k$ 以上が $[10,10,11,12,12,13]$ で $X=4$ の場合、$12,13$ をバラすのがよい

また、$A_i \lt k$ である棒は、最初から長さ $1$ に分割してしまってよい。

実装

外側の二分探索で $O(\log{\max(A)})$ かかるので、判定自体は $O(N \log{何か})$ 程度でおこないたい。

②の判定のために「全ての棒を、次に分割したら $k$ を下回るまで分割する」操作において、実際に棒の長さを配列で持っていたら配列長が $O(N\cdot\max(A))$ になるので、TLE or MLE する。

ただ、$A_i$ から同じ回数だけ分割された棒の長さは、ある整数 $l$ に対し「$l$ または $l+1$」になるので、まとめて管理すれば $O(N)$ で情報を持てる。

（例）$37$ は、1回分割で $(18:1本, 19:1本)$、2回分割で $(9:3本, 10:1本)$、3回分割で $(4:3本, 5:5本)$ となる。

よって、$(l,lの本数,l+1の本数)$ のようにして持っておけば $O(N \log{\max(A)})$ で②の判定に必要な情報を作れる。

各 $A_i$ を分割していく際、終了条件「次に分割したら $k$ を下回る」がややこしく、「$l$ も $l+1$ も下回る」「$l$ は下回るが $l+1$ は $k$ ちょうどになる」を区別して、特に後者は $l+1$ の方だけ②の判定に計上する必要がある。

Python3

def solve(n, m, aaa):
    # k 以上の値が NEEDS 以上作れるなら、中央値は k 以上
    NEEDS = (n + m) // 2 + 1

    def judge(k):

        if sum(a >= k for a in aaa) + m < NEEDS:
            return False

        upper = []
        upper_cnt = 0
        split_cnt = 0
        for a in aaa:
            if a < k:
                split_cnt += a - 1
                continue

            cnt0 = 1
            cnt1 = 0
            for i in range(30):
                b = a >> 1

                if b < k - 1:
                    upper.append((a, cnt0))
                    upper_cnt += cnt0
                    if cnt1 > 0:
                        upper.append((a + 1, cnt1))
                        upper_cnt += cnt1
                    break

                if b == k - 1:
                    if a & 1 == 0:
                        upper.append((a, cnt0))
                        upper_cnt += cnt0
                        if cnt1 > 0:
                            upper.append((k, cnt1))
                            upper_cnt += cnt1
                            split_cnt += cnt1
                            split_cnt += (b - 1) * cnt1
                    else:
                        upper.append((k, cnt0))
                        upper_cnt += cnt0
                        split_cnt += cnt0
                        split_cnt += (b - 1) * cnt0
                        if cnt1 > 0:
                            upper.append((k, cnt1 * 2))
                            upper_cnt += cnt1 * 2
                            split_cnt += cnt1
                    break

                split_cnt += cnt0 + cnt1
                if a & 1 == 0:
                    cnt0 += cnt0 + cnt1
                else:
                    cnt1 += cnt0 + cnt1
                a >>= 1

        # print(f'{n=} {m=} {aaa=} {k=} {upper=} {upper_cnt=} {upper_split=} {lower_split=}')

        if upper_cnt < NEEDS:
            return False

        upper.sort()
        satisfied = False
        tmp = 0
        for val, cnt in upper:
            if satisfied:
                split_cnt += (val - 1) * cnt
                continue
            if tmp + cnt <= NEEDS:
                tmp += cnt
                continue
            satisfied = True
            split_cnt += (val - 1) * (tmp + cnt - NEEDS)

        if split_cnt < m:
            return False

        return True

    l = 0
    r = 1 << 30
    while l + 1 < r:
        mid = (l + r) // 2
        if judge(mid):
            l = mid
        else:
            r = mid
    return l


t = int(input())
for _ in range(t):
    n, m = list(map(int, input().split()))
    aaa = list(map(int, input().split()))
    ans = solve(n, m, aaa)
    print(ans)

目次

AtCoder Beginner Contest 444 E,F問題メモ

E - Sparse Range

問題文

制約

解法

F - Half and Median

問題文

制約

解法

判定方法

実装