ARC 086
C - Not so Diverse
問題
- $N$ 個のボール、各ボールには数字 $a_i$ が書かれている
- いくつかの数字を書き換えることで、数字の種類を $K$ 種類以下にする
- 書き換える数を最小化
解法
逆に書き換えない数字を最大化することを考える。
各数字毎に、書かれたボールの数を数えて、多い方から $K$ 番目までの数が書かれたボールは書き換えずに残しておくと、最大となる。
10 2
1 1 1 2 2 3 4 4 4 5
数字毎にボールの数を数える {1:3, 2:2, 3:1, 4:3, 5:1}
ボール数が多い順にソート [(1,3), (4,3), (2,2), (3,1), (5,1)]
多い方からK=2番目までの数字は 1 と 4
それらが書かれたボールの数は計 6 個で、これをこのまま書き換えないのが最善
残りの4個を書き換えればよい
from collections import Counter
from operator import itemgetter
n, k = map(int, input().split())
a = Counter(map(int, input().split()))
sa = sorted(a.items(), key=itemgetter(1), reverse=True)
if len(sa) <= k:
print(0)
else:
print(sum(c for m, c in sa[k:]))
D - Non-decreasing
問題
- $N$ 要素の数列 $a={a_1,a_2,...,a_N}$
- 「任意の数 $x,y(1 \le x,y \le N)$ を選んで、$a_y$ に $a_x$ を加算」という操作を $2N$ 回まで行える
- 操作の結果、$a$ を単調増加列 $(a_1 \le a_2 \le ... \le a_N)$ としたい
- 操作手順の一例を示せ
解法
$a$ が全て非負なら、累積和を取ると条件を満たす（$N-1$ 回の操作）
$a$ が全て非正なら、後ろから累積和を取ると条件を満たす（$N-1$ 回の操作）
$a$ が正負混在なら、最大値（＝正）と最小値（＝負）、どちらの絶対値が大きいか調べ、最大値が大きければ負の要素に最大値を、最小値が大きければ正の数に最小値を足すと、全て非負または非正にできる（最大$N-1$ 回の操作）
反省
考えれば当たり前なんだけど、なんか累積和がパッと思いつかなかったなあ……
from itertools import starmap
def solve(n, a):
max_a, min_a = float('-inf'), float('inf')
max_i, min_i = 0, 0
for i, b in enumerate(a):
if b > max_a:
max_i, max_a = i + 1, b
if b < min_a:
min_i, min_a = i + 1, b
if max_a >= 0 and min_a >= 0:
return [(i, i + 1) for i in range(1, n)]
if max_a <= 0 and min_a <= 0:
return [(i, i - 1) for i in range(n, 1, -1)]
buf = []
if max_a >= -min_a:
for i, b in enumerate(a):
if b < 0:
buf.append((max_i, i + 1))
buf.extend((i, i + 1) for i in range(1, n))
else:
for i, b in enumerate(a):
if b > 0:
buf.append((min_i, i + 1))
buf.extend((i, i - 1) for i in range(n, 1, -1))
return buf
n = int(input())
a = list(map(int, input().split()))
buf = solve(n, a)
print(len(buf))
print(''.join(starmap('{} {}\n'.format, buf)))
E - Smuggling Marbles
問題
上記リンク参照
解法
根0からの深さが等しいものしか干渉しないので、まずはノードを深さごとに仕分けする。
（※干渉＝初期状態でノード $n_1$ と $n_2$ に置かれたボールが、どこかのタイミングで1つのノードに同時に存在し取り除かれる可能性がある場合、$n_1$ と $n_2$ は干渉するということにする）
また、任意の部分木を考えると、その部分木の中では、外のボールが干渉することは無い。小さな部分木から解いていって、親に遷移するときに兄弟とマージしていく動的計画法が考えられる。
dp[n][i][k] = i 回の操作後に、ノード n に上ってくるボールが k 個になる場合の数（同階層内のみでの場合の数、他の階層を考慮しない）
ただし i = 0は初期状態を示す。また k = 0,1,2以上 の3通りのみ保持。場合の数において他の階層を考慮しない理由は、最後にまとめて計算できるので、途中で逐一やる必要が無いから。
このdpを、階層が深い位置にある $n$ から埋めていく。
最終的に $\sum_{i=0}dp[0][i][1]$ が答え…ではなく、これは違う階層に属するノードの初期状態の組み合わせを考慮できてないので、その組み合わせの数をかける必要がある。
つまり、第 $i$ 階層に属さないノードの数を$m_i$とすると、$\sum_{i=0}(dp[0][i][1] \times 2^{m_i})$ が答えとなる。
０ ／＼ １ ４ ／＼ ＼ ２ ３ ５
dpのやり方: 子が無いノードについて
まず、子が無いノードは、自身に初期状態でボールが置かれてるか置かれてないかしかない。i = 0 のみであり、その時の場合の数[0個,1個,2個以上] = [1,1,0] となる。
よって、ノード 2,3,5について、dp[2] = dp[3] = dp[5] = [[1,1,0]] となる。
dpのやり方: 子があるノードについて
まずi=0は、子が無いノードと同様[1,1,0]である。
i=1以降、子からのボールを統合することになる。2つずつマージできる。
ノード $v$ の子を $c_1,c_2,...$ とし
$$ dp[c_1][i] = [a_0,a_1,a_2] \\ dp[c_2][i] = [b_0,b_1,b_2] $$
であったとき、$c_1$ と $c_2$ の統合結果は、以下になる。（※処理１）
便宜的に、統合の途中結果を $dp[c_{tmp}]$ と表す。（実際にdpへの登録を行うわけでは無く、あくまで $dp[c_{tmp}]$ のデータ構造や意味が $dp[c_1]$ などと同じであることを示す）
\begin{eqnarray} dp[c_{tmp}][i][0] &=& a_0 b_0 \\ dp[c_{tmp}][i][1] &=& a_1 b_0 + a_0 b_1 \\ dp[c_{tmp}][i][2] &=& a_2(b_0 + b_1 + b_2) + a_1(b_1 + b_2) + a_0 b_2 \end{eqnarray}
数式で書くとわかりにくいが、要は、以下のような場合の数の計算をしているだけ
- (統合結果が1) = (c1から1 ∩ c2から0) ∪ (c1から0 ∩ c2から1)
- (統合結果が1の場合の数) = (c1が1の場合の数 × c2が0の場合の数) ＋ (c1が0の場合の数 × c2が1の場合の数)
2つの子の深さが異なるとき、例えば $c_2$ より $c_1$ の方が部分木が深く、$dp[c_1][i]$ に対応する $dp[c_2][i]$ が無い、という時は、それ以降の $i$ については $dp[c_1][i]$ のままでよい
順次 $c_{tmp}$ にマージしていき、全ての子について統合が完了すると、次は「同じノードに2個以上同時にボールがあると取り除く」という操作に基づき、各 $i$ につき、以下を行う。（※処理２）
\begin{eqnarray} dp[c_{tmp}][i][0] &+=& dp[c_{tmp}][i][2] \\ dp[c_{tmp}][i][2] &=& 0 \end{eqnarray}
この処理の後、$dp[c_{tmp}]$ の先頭に[1,1,0]を付け加えたものが、$dp[v]$ となる。
ただ、これをそのまま実装すると、部分点は取れるが全ては時間的に間に合わないので、不要な処理をカットする。
子の数が1つの時
マージする必要は無く、$dp[v]$ は $dp[c]$ の先頭に[1,1,0]を加えればよい
子の数が2つ以上の時
最も深い子の深さ（dp[c]の長さ）を $d_1$, 2番目に深い子の深さを $d_2$ とすると、上記の※処理２（統合後、2個以上の場合の数を0個に移す処理）は、$d_2 \le i \lt d_1$の範囲については、行う必要は無い。
dpに登録するときは必ず※処理２によって「2個以上」の場合の数を 0 にしているので、更新が無かった部分については繰り返し行う必要が無い。
これらは、特に階層の深い入力ケースに対する枝刈りとして有効。
データ構造について
子を統合するとき、新しく配列を作るのでは無く「dp[c]の中で最長のもの」に対して統合するようにすると、余計なメモリを消費せずに済み、処理時間も速くなる。
その際、[1,1,0]を先頭に付け加える処理を頻繁に行うので、各 dp[n] は双方向リストで持つとよい。
仮に全てのdp[n]を保存していた場合、$N \le 2 \times 10^5$ で、全てが直列に繋がっているような意地悪なケースだと、dp[n]のサイズは下からの階層に比例するので上は $N$ に近くなり、それが $N$ 個あり、それぞれが3つの状態を記録するので、$3N(N-1)/2 = 6 \times 10^{10}$ 個のintをもつことになり、MLEの可能性もある。
処理の終わった階層の情報が正しくなくなるが、問題を解くには使わないので問題ないことにする。
from collections import deque
def get_pow():
cache = {}
def func(x):
if x not in cache:
cache[x] = pow(2, x, mod)
return cache[x]
return func
mod = 1000000007
n = int(input())
parents = list(map(int, input().split()))
children = [set() for _ in range(n + 1)]
for c, p in enumerate(parents):
children[p].add(c + 1)
levels = [{0}]
while True:
level = set()
for p in levels[-1]:
level.update(children[p])
if not level:
break
levels.append(level)
levels.reverse()
level_node_count = []
balls = [None] * (n + 1)
for i, level in enumerate(levels):
level_node_count.append(len(level))
for node in level:
cn = children[node]
if cn:
if len(cn) == 1:
bs = balls[next(iter(cn))]
bs.appendleft([1, 1, 0])
balls[node] = bs
continue
balls_from_children = [balls[c] for c in children[node]]
balls_from_children.sort(key=len)
bs1 = balls_from_children[0]
for bs2 in balls_from_children[1:]:
for (b10, b11, b12), b2 in zip(bs1, bs2):
b2[2] = ((b11 + b12) * b2[1] + b12 * b2[0]) % mod
b2[1] = (b10 * b2[1] + b11 * b2[0]) % mod
b2[0] = b2[0] * b10 % mod
bs1 = bs2
lim = len(balls_from_children[-2])
for i, b in enumerate(bs1):
if i >= lim:
break
b[0] = (b[0] + b[2]) % mod
b[2] = 0
bs1.appendleft([1, 1, 0])
balls[node] = bs1
else:
balls[node] = deque([[1, 1, 0]])
level_node_count.reverse()
pow2 = get_pow()
print(sum(b[1] * pow2(n - l + 1) % mod for l, b in zip(level_node_count, balls[0])) % mod)
