AtCoder Beginner Contest 420 F,G問題メモ

AtCoder Beginner Contest 420

珍しく日曜日の開催。

F - kirinuki

• $N \times M$ のグリッドが与えられます。各マスには '.' または '#' が書かれています。
• グリッド内に、以下をともに満たす長方形領域はいくつありますか？
  • '.' が書かれているマスのみからなる
  • 面積は $K$ 以下

• $N,M,K$ は整数
• $1 \le N,M \le 5 \times 10^5$
• $1 \le N \times M \le 5 \times 10^6$
• $1 \le K \le N \times M$
• $S_i$ は '.' と '#' からなる長さ $M$ の文字列

最大長方形を求める問題と似ているが、そこを出発点としてもそこそこ考察が必要。

• ヒストグラム内の最大長方形を求めるアルゴリズム
• 競技プログラミングにおける細かな話題まとめ - はまやんはまやんはまやん

「グリッド内の最大長方形」は、長方形に使う下端の行を決め打って「ヒストグラム内の最大長方形」を各行に対して解くことで求められる。
「自身を含め、自身の上にいくつ '.' が連続しているか」がヒストグラムの高さに相当する。
スタックを使うと1行当たり $O(M)$ で求められる。

同様にこの問題も、「各行毎に、そこを下端とする長方形の個数を、ヒストグラムに見立てて数える」ことを考える。

まずいったん、面積が $K$ 以下という条件は忘れて、単純に「ヒストグラム内の長方形の個数」を数える問題を考える。

ある行において、$j$ 列目のヒストグラムの高さを $h_j$ とする。
もし、$h_j$ が全て異なっている場合、以下のように分割統治で答えを求められる。

j   1 2 3 4 5 6 7 8
h   6 3 9 4 7 2 5 8
   
       □
       □        □
       □  □    □
   □  □  □    □
   □  □  □  □□
   □  □□□  □□
   □□□□□  □□
   □□□□□□□□
   □□□□□□□□

まず、最小の h を持つ j=6 について、
j=6 の下端を必ず使い、高さが hj=2 以下の長方形を数える。

       △
       △        △
       △  △    △    ■: 必ず使う
   △  △  △    △    □: 使ってよい
   △  △  △  △△    △: そもそも■を使おうとしたら使えない
   △  △△△  △△
   △△△△△  △△
   □□□□□□□□
   □□□□□■□□
   
他の h は必ず hj より高いので、使ってもよい候補の□は横いっぱいに長方形状に並ぶ。
左側と右側をどこまで使うかで、上例の場合、横幅は 6 x 3 = 18 通りある。
また、高さは hj=2 通りあるので、合計 36 通りの長方形が作れる。

j=6の下端を使わないものは、以降、左右に分割して考えられる。
次に小さい h を持つ j=2 について、
j=2 の下端を必ず使い、高さが hj=3 以下の長方形を数える。

       △
       △          
       △  △          ■: 必ず使う
   △  △  △          □: 使ってよい
   △  △  △          △: そもそも■を使おうとしたら使えない
   △  △△△
   □□□□□
   □□□□□×
   □■□□□×

横幅の取り方は 2 x 4 = 8 通り、高さは hj=3 通りある。
合計 24 通りの長方形が作れる。

       △
       △          
       △  △          ■: 必ず使う
       △  △          □: 使ってよい
       △  △          △: そもそも■を使おうとしたら使えない
       □□□
     ×□□□
     ×□□□×
     ×□■□×

次に小さいのは j=4, hj=4
2 x 2 x 4 = 16 通りの長方形が作れる。

$h_j$ の小さい順にこのように、$(高さ, 左端の選択肢数, 右端の選択肢数)$ によって計算できる。
そしてこれは、$h_j$ に同じものが含まれる場合でも、適当に順序を付けてやると問題なく求まる。

前述のように、ある下端マスを必ず使う長方形の個数は(高さ、左端の選択肢、右端の選択肢)の組 $(h,l,r)$ によって決まる。

面積が $K$ 以下という制限を考慮すると、 横幅が $w$ の場合、高さは $\lfloor \frac{K}{w} \rfloor$ 以下という制約が加わる。
よって、横幅毎にパターン数を分類して考えたい。

   △  △  △    △△
   △△△  △△△△△    高さ h = 4
   □□□□□□□□□    左端の選択肢 l = 4
   □□□□□□□□□    右端の選択肢 r = 6
   □□□□□□□□□
 ×□□□■□□□□□×

横幅 1
 ×□□ [■] □□□□×    1通り
 
横幅 2
 ×□ [□■] □□□□×    2通り
 ×□□ [■□] □□□×

横幅 3
 × [□□■] □□□□×    3通り
 ×□ [□■□] □□□×
 ×□□ [■□□] □□×

横幅 4
  [□□□■] □□□□×    4通り
 × [□□■□] □□□×
 ×□ [□■□□] □□×
 ×□□ [■□□□] □×

横幅 5
  [□□□■□] □□□×    4通り
 × [□□■□□] □□×
 ×□ [□■□□□] □×
 ×□□ [■□□□□] ×

横幅 6
  [□□□■□□] □□×    4通り
 × [□□■□□□] □×
 ×□ [□■□□□□] ×
 ×□□ [■□□□□□] 

横幅 7
  [□□□■□□□] □×    3通り
 × [□□■□□□□] ×
 ×□ [□■□□□□□]

横幅 8
  [□□□■□□□□] ×    2通り
 × [□□■□□□□□]

横幅 9
  [□□□■□□□□□]     1通り

横幅 $w$ に対するパターン数 $p(w)$ は以下のような形となることが分かる。対称性より、$l \le r$ を仮定する。

• $1 \le w \le l$ の範囲は1ずつ増える（$p(w)=w$）
• $l \lt w \le r$ の範囲は $l$ 固定（$p(w)=l$）
• $r \lt w \le l+r-1$ の範囲は1ずつ減る（$p(w)=l+r-w$）

また横幅 $w$ に対しては、高さの選択肢が $\min(h, \lfloor \frac{K}{w} \rfloor)$ 通りある。
これを $p(w)$ と掛け合わせて、総和を取ったものが答えとなる。

w  p(w) min(h,K/w)          K=15, h=4
1   1   x   4   =   4  ┐
2   2   x   4   =   8  ┤
3   3   x   4   =  12  ┤
4   4   x   3   =  12  ┤
5   4   x   3   =  12  ┼ 65個
6   4   x   2   =   8  ┤
7   3   x   2   =   6  ┤
8   2   x   1   =   2  ┤
9   1   x   1   =   1  ┘

これを愚直に計算していたら $O((NM)^{1.5})$ となりTLEなので、高速化を考える。

$p(w)$ は $l,r$ によって3つの区間に分かれ、それぞれ規則的なのでまとめて考えられそう。

また、$h \le \lfloor \frac{K}{w} \rfloor$ となる最大の $w$ を $x$ とする。
$w$ が $x$ 以下か、より大きいかで、$p(w)$ にかかる係数が $h$ か $K/w$ か分かれるので、これも2つに計算を分けて考える。

つまり「$p(w)$ のどの区間か」「$x$ より上か下か」で、 計6通りの場合について、区間の計算を $O(1)$ でできるように前計算しておけばよい。

p(w)
 ^        _____
 |      ／|   |＼
 |    ／  |   |  ＼
 |  ／    |   |    ＼
 |／      |   |      ＼
 +--------|---|----------> w    [1,3]: (1ずつ増える) × 固定値 h
   1 2 3 4 5 6  7 8 9           [4,4]: (1ずつ増える) × K/w
         l   r                  [5,6]:    固定値 l   × K/w
                                [7,9]: (1ずつ減る)   × K/w
min(h,K/w)                             ~~~~~~~~~~~~~    ~~~~~~~~
 ^ _____                                   3通り          2通り
 |      \___
 |      |   \___                l,r,x によって分割されたこの4つの区間を
 |      |       \___            それぞれ O(1) で求められればよい
 |      |           \
 +------|--------------> w
   1 2 3 4 5 6 7 8 9
       x

以下を前計算しておく。（$A,B$ は累積和を取ることで間接的に求められる）

• $\displaystyle A_{l,r}:=\sum_{w=l}^{r} \lfloor \frac{K}{w} \rfloor$
• $\displaystyle B_{l,r}:=\sum_{w=l}^{r} w \times \lfloor \frac{K}{w} \rfloor$
• $C_{h}:=h$ に対する $x$（※ $h \le \lfloor \frac{K}{w} \rfloor$ となる最大の $w$）

区間 $[a,b]$ における答え $\displaystyle \sum_{w=a}^b p(w) \min(h,\lfloor \frac{K}{w} \rfloor)$ は、 $l,r,x$ との大小関係によって、以下のように求められる。

• $a \le b \le l, a \le b \le x$ （1ずつ増える × $h$）
  • $\frac{(a+b)(b-a+1)}{2} \times h$
• $a \le b \le l, x \lt a \le b$ （1ずつ増える × $K/w$）
  • $B_{a,b}$
• $l \lt a \le b \le r, a \le b \le x$ （$l$ × $h$）
  • $l \times h \times (b-a+1)$
• $l \lt a \le b \le r, x \lt a \le b$ （$l$ × $K/w$）
  • $l \times A_{a,b}$
• $r \lt a \le b \lt l+r, a \le b \le x$ （1ずつ減る × $h$）
  • $\frac{(l+r-a+l+r-b)(b-a+1)}{2} \times h$
• $r \lt a \le b \lt l+r, x \lt a \le b$ （1ずつ減る × $K/w$）
  • $(l+r) \times A_{a,b} - B_{a,b}$

$l,r,x$ によって分割される最大4区間のそれぞれを上記のうちの適切な式でまとめて計算すると、 $(h,l,r)$ に対する答えを $O(1)$ で求められる。

$h_j,A_{l,r},B_{l,r}$ などは、前計算しておく。

ヒストグラム中の最大長方形のアルゴリズムと同じように、スタックを用いて管理していく。

• Cartesian Tree の実装 〜 stack を用いた O(N) 時間構築 - けんちょんの競プロ精進記録

スタックに $(j,h_j)$ を積もうとする時、既にスタックに積まれている要素のうち、$h_k \ge h_j$ となる $(k,h_k)$ はpopされる。
$k$ についての長方形数の計算は、この「popされたタイミング」でおこなう。

[(-1,-1), (3,5), (6,6), ]  ←  (9,4)

スタックが上のような状況で、$(j,h_j)=(9,4)$ を積もうとして、$(k,h_k)=(6,6)$ をpopしたとする。
$(6,6)$ のさらに左にある要素を $(m,h_m)=(3,5)$ とする。

この時、以下の図のような状態になっているので、

   △△  △△
   □□□□□      ■: 必ず含める
 ×□□□□□      □: 含めてもよい
 ×□□□□□×    △: 含めることができない
 ×□□□□□×    ×: 含めてはいけない
 ×□□□□□×
 ×□□■□□×
  m     k     j
j 3 4 5 6 7 8 9
h 5     6     4

$k$ の下端を含む長方形の個数は、$(h,l,r)=(h_k, k-m, j-k)$ として求められる。

Python3

def solve(n, m, k, field):
    space_above = [[0] * (m + 1) for _ in range(n)]
    for j in range(m):
        space_above[0][j] = 1 if field[0][j] == '.' else 0
    for i in range(1, n):
        for j in range(m):
            space_above[i][j] = space_above[i - 1][j] + 1 if field[i][j] == '.' else 0

    div_k = [0] * (m + 2)
    div_acc = [0] * (m + 2)
    rev_k = [0] * (k + 1)
    for i in range(1, m + 2):
        div_k[i] = k // i
        div_acc[i] = div_acc[i - 1] + k // i
    for i in range(1, k + 1):
        rev_k[min(k // i, n)] = i
    for i in range(k - 1, -1, -1):
        rev_k[i] = max(rev_k[i], rev_k[i + 1])

    mul_acc = [0] * (m + 2)
    for i in range(1, m + 2):
        mul_acc[i] = mul_acc[i - 1] + i * div_k[i]

    def get_ans(h, l, r):
        if l > r:
            l, r = r, l
        res = 0
        x = min(rev_k[h], l + r) if h <= k else 0
        if x <= l:
            res += h * (1 + x) * x // 2
            res += mul_acc[l] - mul_acc[x]
            res += l * (div_acc[r] - div_acc[l])
            res += (l + r) * (div_acc[l + r] - div_acc[r]) - (mul_acc[l + r] - mul_acc[r])
        else:
            res += h * (1 + l) * l // 2
            if x <= r:
                res += h * l * (x - l)
                res += l * (div_acc[r] - div_acc[x])
                res += (l + r) * (div_acc[l + r] - div_acc[r]) - (mul_acc[l + r] - mul_acc[r])
            else:
                res += h * l * (r - l)
                res += h * (l - 1 + l - (x - r)) * (x - r) // 2
                res += (l + r) * (div_acc[l + r] - div_acc[x]) - (mul_acc[l + r] - mul_acc[x])
        return res

    def get_ans_naive(h, l, r):
        res = 0
        for i in range(l):
            for j in range(r):
                for t in range(1, h + 1):
                    w = i + j + 1
                    if w * t <= k:
                        res += 1
        return res

    ans = 0
    # ans_naive = 0
    for i in range(n):
        stack = [(-1, -1)]
        for j in range(m + 1):
            sa = space_above[i][j]
            while stack[-1][1] >= sa:
                p, h = stack.pop()
                q = stack[-1][0]
                l = p - q
                r = j - p
                ans += get_ans(h, l, r)
                # ans_naive += get_ans_naive(h, l, r)
                # print(f'{i=} {p=} {h=} {q=} {l=} {r=} {get_ans(h, l, r)=} {get_ans_naive(h,l,r)=}')
            stack.append((j, sa))
    # print(f'{ans_naive=}')
    return ans


n, m, k = map(int, input().split())
field = [input() for _ in range(n)]
ans = solve(n, m, k, field)
print(ans)

G - sqrt(n²+n+X)

• 整数 $X$ が与えられます。
• $\displaystyle \sqrt{n^2+n+X}$ が整数となるような整数 $n$ を全て求めてください。

• $-10^{14} \le X\le 10^{14}$
• 入力される値は整数

FとGのAC人数がここまで大幅に逆転するのも珍しい。
どうも受験数学における典型問題らしいこと、 また制約から何らかの $O(\sqrt{X})$ 解法であることは想定つき、 未証明でも勘で解法を当てやすい問題ではあったことなどで、むべなるかな。

公式Editorialの解は、$n^2+n+X=m^2$ とおくと $(2m-2n-1)(2m+2n+1)=4X-1$ と式変形できるので、 $4X-1$ の約数を列挙して連立方程式を解いて整数になる $n$ を求める、というもの。

$n$ についての平方完成を挟む式変形が、慣れていないとパッとは思いつきづらい。

$n^2+n+X=(n+m)^2$ とおくと $n=\frac{X-m^2}{2m-1}$ と表せる。
$m$ を $-10^7～10^7$ くらいまで調べることでもACできる。

ただ、$m$ の範囲がこれでいい証明はどうすれば？

• 解説 - AtCoder Beginner Contest 420

上のユーザー解説のように $4n=\dfrac{4X-1}{2m-1}-(2m+1)$ と表せることから、 $4X-1$ の約数の全列挙で解けるのはわかるのだが、その中には $m \ge \sqrt{X}$ となる $m$ もあるはず。
でも、「$m$ を $-10^7～10^7$ の範囲で試す」だけの列挙で、ACできてしまってるんだよな……。

• 追記: こうすると証明できる（nesya氏ユーザー解説）
  • 解説 - AtCoder Beginner Contest 420

自分の本番時の解法。上記よりもう少し（無駄に）こねくり回している。

$n^2+n+X=(n+m)^2$ とおき、$m$ について整理すると $m^2+2nm-(n+X)=0$ となり、これが整数解を持つ。
2つの解を $a,b$ とすると、$a+b=2n, ab=-(n+X)$ から $b=-\dfrac{2X+a}{2a+1}$ という関係が得られる。

これは横軸に $a$、縦軸に $b$ をとったグラフで表すと、 $(-\dfrac{1}{2},-\dfrac{1}{2})$ を中心とした反比例の形になっている。

$X \gt 0$ の時、双曲線は第2象限と第4象限に広がる。
この2つは $a=b$ の直線に対し対称的（$a,b$ を入れ替えただけ）なので、第4象限の部分のみに限定して考える。
このグラフが通る格子点で、最も左にあるのは $(a,b)=(0,-2X)$ である。最も右にあるのは、$(2X-1,-1)$ である。

$a=0～\sqrt{2X}$ の範囲は $a$、 それ以降は $b=-\sqrt{2X}-1～-1$ の範囲の $b$ を整数として全探索することで、 格子点を全て調べられ、そこから $n$ も求められる。

$X \lt 0$ の時も、少しずれているものの、同様に考えられる。

Python3

def sqrt(x):
    res = int(x ** 0.5)
    while (res + 1) ** 2 <= x:
        res += 1
    while res ** 2 > x:
        res -= 1
    return res


def solve(x):
    eps = 100
    ans = []
    for a in range(sqrt(abs(x)) + eps):
        num = a + x * 2
        div = 2 * a + 1
        if num % div != 0:
            continue
        b = -num // div
        assert (a + b) % 2 == 0
        n = (a + b) // 2
        ans.append(n)

        a = -a - 1
        num = a + x * 2
        div = 2 * a + 1
        if num % div != 0:
            continue
        b = -num // div
        assert (a + b) % 2 == 0
        n = (a + b) // 2
        ans.append(n)

    ans = set(ans)
    ans = sorted(ans)
    return ans


x = int(input())
ans = solve(x)
print(len(ans))
print(*ans)