AtCoder Beginner Contest 159 D,E,F問題メモ

AtCoder Beginner Contest 159

D - Banned K

問題

$N$ 個のボールが一列に並び、 $i$ 番目のボールに書かれた整数は $A_{i}$
書かれた整数が同じ2個のボールを取り出す方法の個数を考える
各 $1 \leq k \leq N$ について、 $k$ 番目のボールを使わないで取り出す方法の個数を答えよ
$3 \leq N \leq 2 \times 10^{5}$

解法

同じ整数が書かれたボールが $C$ 個あったら、そこから2つを選ぶので、 $\frac{C (C - 1)}{2}$ 通り。

これを各整数について求めて足し合わせれば、 $k$ 番目のボールを使わないという条件を無視した答えが得られる。

$k$ 番目のボールを使わない場合、その整数だけ $C - 1$ に置き換えて $\frac{(C - 1) (C - 2)}{2}$ とすればよい。

ただし、 $k = 1 ～ N$ について他の整数のパターン数をその都度また計算して足し合わせていては時間がかかるので、全てのボールを使う場合から、差分だけ更新する。

$k$ 番目のボールに書かれた整数が、 $c$ 個ある整数だったなら、 $- \frac{c (c - 1)}{2} + \frac{(c - 1) (c - 2)}{2}$ とすればよい。

 
 
 
 
 from collections import Counter
 
n = int(input())
aaa = list(map(int, input().split()))
cnt = Counter(aaa)
 
base = 0
for c in cnt.values():
    base += c * (c - 1) // 2
 
buf = []
cache = {}
for a in aaa:
    if a in cache:
        buf.append(cache[a])
        continue
    c = cnt[a]
    ans = base - c * (c - 1) // 2 + (c - 1) * (c - 2) // 2
    buf.append(ans)
    cache[a] = ans
print('\n'.join(map(str, buf)))

 

 
 

E - Dividing Chocolate

問題

$H \times W$ のグリッド状の板チョコレートがある
このチョコはグリッド毎に普通のチョコとホワイトチョコが混ざっている
- $S_{i, j} = 0$ の場合、上から $i$ 行目、左から $j$ 列目のグリッドは普通のチョコ
- $S_{i, j} = 1$ の場合、上から $i$ 行目、左から $j$ 列目のグリッドはホワイトチョコ
このチョコをグリッドの境界で切って、1つの欠片に含まれるホワイトチョコを、全ての欠片で $K$ マス以下にしたい
切るときは、既にいくら切られていようと「元の状態での全体」から縦または横一直線に切る（問題ページ参照）
最小の切る回数を求めよ
$1 \leq H \leq 10$
$1 \leq W \leq 1000$
$1 \leq K \leq H \times W$

解法

わりと実装寄りの問題。

$H$ が異様に小さくて切る位置の組み合わせをbit全探索してくださいというメッセージを受信したのでそうする。たかだか $2^{9}$ 。

縦が決まれば、後は左端から、全てのブロックが $K$ を越えないギリギリの位置で切っていく、ということを繰り返せばよい。

K=3
  1 1 1 0 1 1 1
  1 0 0 1 1 0 1
✂-------------
  0 1 1 1 1 1 1
  0 0 1 1 0 0 1
     ^ ^   ^

二次元累積和を使ったら速いかなと思ったけど、添え字の変化を追うのが難しくなった。

解説放送でやっているように、各行が何番目のグループに属するかを求めて、グループ毎にホワイトチョコの個数をカウントする配列を用意して、1列ずつ加算していった方がやりやすかったかも。

やることは分かっても、スッキリする実装ができるかどうか、安定しないね。

初回通した実装

（ちょっとは）意味を追いやすい実装

 
 
 
 
 def column_check(grp, grp_count, choco, h, j, k):
    """ 0:成功（更新も反映）  1:失敗（区切った結果を返す） 2:不可能 """
    # この行のグループ毎のカウント
    curr_count = [0] * len(grp_count)
    # この行だけを数える
    for i in range(h):
        if choco[i][j] == '0':
            continue
        curr_count[grp[i]] += 1
    # この行だけでk越えてたら現在の横の切り方では不可能
    for cc in curr_count:
        if cc > k:
            return 2
    # ここまでの行と合わせてkを越えないかチェック
    for i, cc in enumerate(curr_count):
        grp_count[i] += cc
        # 越えてたらその直前で切るので、grp_countを現在の行に上書き
        if grp_count[i] > k:
            grp_count[:] = curr_count
            return 1
    return 0
 
 
def solve(h, w, k, choco):
    ans = 10 ** 9
    grp = [0] * h
    for bit in range(1 << (h - 1)):
        g = 0
        for i in range(h):
            grp[i] = g
            if (bit >> i) & 1:
                g += 1
 
        grp_count = [0] * (g + 1)
        tmp = 0
        possible = True
        for j in range(w):
            result = column_check(grp, grp_count, choco, h, j, k)
            if result == 2:
                possible = False
                break
            if result == 1:
                tmp += 1
 
        if not possible:
            continue
 
        ans = min(ans, tmp + bin(bit).count('1'))
 
    return ans
 
 
h, w, k = map(int, input().split())
choco = [input() for _ in range(h)]
print(solve(h, w, k, choco))

 

 
 

F - Knapsack for All Segments

問題

長さ $N$ の数列 $A_{1}, . . ., A_{N}$ と整数 $S$ が与えられる
$f (L, R)$ （ $1 \leq L \leq R \leq N$ ）を以下で定義する
- $A_{L} ～ A_{R}$ の範囲にある数の組み合わせ（何個でもいい）で、合計がちょうど $S$ のものの数
ここで、 $L, R$ の取り方は $\frac{N (N + 1)}{2}$ 個あるが、全てについての $f (L, R)$ の総和を求めよ
$1 \leq N, S, A_{i} \leq 3000$

解法

解く上で主眼を「 $(L, R)$ の組」から「合計が $S$ になる組」に移す。合計が $S$ になる組をすっぽり包む $(L, R)$ を数えて合計すれば、結局答えは同じことになる。

S = 9

i   1  2  3  4  5
A   3  1  4  1  5    （便宜上、1を区別するのに1a,1bとする）

(L,R)の組       範囲内に含まれる合計=Sとなる整数の組
[1, 4]          {3, 1a, 4, 1b}
[1, 5]          {3, 1a, 5}, {3, 1b, 5}, {3, 1a, 4, 1b}, {4, 5}
[2, 5]          {4, 5}
[3, 5]          {4, 5}
※他の組は0
                            ⇅個数は一致

合計=Sの組      この組が数え挙げられる(L,R)の組
{3, 1a, 5}      [1, 5]
{3, 1b, 5}      [1, 5]
{3, 1a, 4, 1b}  [1, 4], [1, 5]
{4, 5}          [1, 5], [2, 5], [3, 5]

DPで考えるとして、区間の右端 $R$ を決め打ってみる。

以下のように、 $i$ を1つ進めるたび、添え字を $A_{i}$ だけずらして元の配列に足し合わせていくことをまず考えつく。

i    (0)   1    2    3    4    5
A          3    1    4    1    5

DP 0  1    1    1    1    1    1
   1            1    1    2    2
   2                      1    1
   3       1    1    1    1    1
   4            1    2    3    3
   5                 1    3    4
   6                      1    3
   7                 1    1    2
   8                 1    2    3
   9                      1    4

しかし、これで数えられるのは、単に合計が $S$ になる組の個数である。

つまり、例えば {4, 5} も1個として数えられている。今回は、この組の場合は左端の位置を1～3まで選べるので、これは3個として数えたい。

これを実現するにはどうするか。

要は、各組の最も左の要素が $i$ であれば、区間の左端は $1 ～ i$ までを選べるということなので、はじめて $D P [0] \to D P [それ以外]$ に遷移するタイミングで個数を調整してやればよい。一度左端が決まった後の遷移は、変わらない。

i    (0)   1    2    3    4    5
A          3    1    4    1    5

DP 0  0    1    2    3    4    5    ←DP[0]を1ずつ増やす（区間の左端を選択できる個数）
   1            2    2    6    6
   2                      2    2
   3       1    1    1    1    1
   4            1    4    5    5
   5                 2    6   11
   6                      2    8
   7                 1    1    3
   8                 1    2    3
   9                      1    6

後は、各 $i$ につき $D P [S]$ を合計していけばよい。

「各組の最も右の要素が $i$ であれば、区間の右端は $i ～ N$ まで選べる」という、左を反転した性質については、 $i$ が進んでもDP配列を次に引き継ぎ、その都度答えに足し合わせていく、という処理によって組み込まれている。

左右、対称性のある性質なのに、数え上げの処理が異なる、というのが何とも奇妙な感じ（そういう風な解法を取ったからだが）。

 
 import numpy as np
 
n, s = map(int, input().split())
aaa = list(map(int, input().split()))
 
dp = np.zeros(s + 1, dtype=np.int64)
 
ans = 0
MOD = 998244353
for a in aaa:
    dp[0] += 1
    if a <= s:
        dp[a:] = (dp[a:] + dp[:-a]) % MOD
    ans = (ans + dp[s]) % MOD
print(ans)

なお、AtCoderの現在のNumPyのバージョン（1.8.2）では、配列に自身をずらした配列を加算代入する処理にバグがあるため注意。

×:
aaa[3:] += aaa[:-3]

○:
aaa[3:] = aaa[3:] + aaa[:-3]
aaa[3:] += aaa[:-3].copy()