AtCoder Grand Contest 046 B,C,D,E問題メモ

AtCoder Grand Contest 046

YouTube解説はありません。

B - Extension

問題

  • 縦横 $A \times B$ のマス目があり、全てのマスは白く塗られている
  • 以下の操作を、マス目が $C \times D$ になるよう繰り返し行う
  • 操作
    • 縦または横を選ぶ
      • 縦を選んだ場合、マス目の上に1行追加する
      • 横を選んだ場合、マス目の右に1列追加する
    • 追加されたマスのうち、1マスを黒く、残りを白く塗る
  • 最終的なマス目の異なる塗られ方としてありうるものの個数を 998244353 で割った余りを求めよ
  • $1 \le A \le C \le 3000$
  • $1 \le B \le D \le 3000$

解法

誤読して、行・列の追加は好きな位置に挿入できると捉えちゃってた。 日常では1マスの意味で「マス目」って使うこともあるけど、厳密には「マスが並んだ模様」を指すらしい。 「網目」が糸と糸の隙間の1つを指すように、「目」ってどこかピンポイントの1つというイメージがある。 でも、確かに「柾目」だと木の模様という意味で使われる。ニホンゴムツカシイ。

ともかく、最上段の上と、再右列の右にしか追加されないなら、2次元DPでいけそう。

縦 $i$ 列、横 $j$ 列のマス目を $(i,j)$ と表すとして、以下のDPを定義する。

  • $DP[i][j]=(i,j)$ の塗り方のパターン数

右に追加するとき、行数だけ黒く塗るマスの選択肢がある。

┌ □□..□○
i  ..    ..○
└ □□..□○

上に追加するとき、列数だけ黒く塗るマスの選択肢がある。

   ○○..○
   □□..□
   ..    ..
   □□..□
   └  j ┘

なので、マス目が $(i,j)$ になった時の塗り方は、$(i-1,j)$ と $(i,j-1)$ から計算できる。

  • $DP[i][j] = DP[i-1][j] \times j + DP[i][j-1] \times i$(暫定)

ただしこれは重複を含んでいて、たとえば以下の2つは同じになる。(○が追加マス)

○○●○    □□■○
□■□□    □■□○
□□■□    □□■○
□□□■    □□□●

どういうケースかというと、最上行と再右列に、共に黒マスが1つずつある状態が該当する。

①①①□    ①②からそれぞれ1つが黒いケースは重複
□■□②
□□■②
□□□②

最も右上のマスが黒い場合は以下の理由により気にしなくて大丈夫。

  • 最上行か再右列に他に黒マスがあるケースは、最後に追加された方向が特定できるので重複しない
  • 最上行と再右列で右上の1マスしか黒くない、という状態はあり得ない
    • たとえば行の追加によりその状態になったなら、その直前に列が追加された際、最右列のどれか1つは黒いはず
    • 逆も同じ
  • 唯一、$(A,B)$ から行か列のどちらかにしか追加していない場合はあり得るが、そこは最初に独自に計算してしまえばよい

重複ケースは $(i-1)(j-1)DP[i-1][j-1]$ 通りあり得るので、これを引いたものが正しい遷移となる。

  • $DP[i][j] = j \times DP[i-1][j] + i \times DP[i][j-1] - (i-1)(j-1)DP[i-1][j-1]$

a, b, c, d = map(int, input().split())
MOD = 998244353

dp = [[0] * (d + 1) for _ in range(c + 1)]
dp[a][b] = 1

for i in range(a + 1, c + 1):
    dp[i][b] = dp[i - 1][b] * b % MOD

for j in range(b + 1, d + 1):
    dp[a][j] = dp[a][j - 1] * a % MOD

for i in range(a + 1, c + 1):
    for j in range(b + 1, d + 1):
        dp[i][j] = (dp[i][j - 1] * i + dp[i - 1][j] * j - dp[i - 1][j - 1] * (i - 1) * (j - 1)) % MOD

print(dp[c][d])

C - Shift

問題

  • '0'と'1'のみからなる文字列 $S$ と、整数 $K$ が与えられる
  • 以下の操作を0回以上 $K$ 回以下繰り返してできる可能性のある文字列の個数を 998244353 で割った余りを求めよ
  • 操作
    • $S_i=$'0'、$S_j=$'1'、$i \lt j$ であるような $(i,j)$ の組を選ぶ
    • $S_j$ を取り除き、$S_i$ の直前に挿入する
  • $1 \le |S| \le 300$
  • $0 \le K \le 10^9$

解法

上手く問題設定を言い換えてDP。

複数の操作手順で同じ最終結果が得られうるので、どう考えたら重複を除けるか。

「最終結果の各'0'の間に、'1'が何個あるか?」を考えるとよい。 位置が変わらないものを固定してその各間の状態を考える、というのは解法としてたまに見る。

また、無駄な操作、というのがあり得るが、

  A  BCD  E    EをDの直前に移動する操作と、
..0..110..1..  BやCをAの直前に移動する操作を同時にやるのは、
               最初からEをAに持っていくのと変わらない

幸い、問題設定は「0回以上 $K$ 回以下」なので、このように2回に分けて1回と同じ操作になるケースは とりあえず最短手数として計算しておけばよい。 (無駄な操作も無限にできるわけではないので、ちょうど $K$ 回という問題設定ならより複雑な気がする)

すると、$K$ の制約は大きいものの、実際には'1'の個数回までしか操作する必要はないので、上限がぐっと小さくなる。

  • $K←\min(K, 1の個数)$

初期状態で $i$ 番目の'0'の前にある'1'の数 $B_i$ を計算する。

S=01100110

     011 0   011 0
B  0   2   0   2   0

1回の操作は、0以上の $B_i$ を1つ減らして、より左のどこかを1増やすことに相当する。

   v---,                    対応するS
   1   1   0   2   0   ...  10100110
   v-----------,
   2   1   0   1   0   ...  11010010

DP

初期 $B_i$ から操作を繰り返した状態を $C_i$ として、最終状態としてあり得る $C_i$ を左から決めていく。

  • $DP[i][j][k]=C_i$ までの値を決めて、移動先の位置が確定した'1'が $j$ 個、移動元の位置が確定した'1'が $k$ 個の場合のパターン数

$DP[i]$ の更新には直前の $DP[i-1]$ しか必要でないため、 $i$ を省略して2次元配列を逐次更新する方法で考える。

$j,k$ の取りうる値の範囲は、移動しうる'1'の最大個数なのでそれぞれ $0~K$ となる。$j \lt k$ の部分はあり得ないので考えなくてよい。

初期状態          B
j\k  0  1  2      0   2   0   2   0
0    1  -  -
1    0  0  -
2    0  0  0

各 $C_i$ では、$j=K$ になるまでそこを増やすことができる。 これは、下方向に累積和を取る操作となる。

i=0               C0   対応する(j, k)
j\k  0  1  2
0    1  -  -      0   (0, 0)
1    1  0  -      1   (1, 0)
2    1  0  0      2   (2, 0)

$B_i \gt 0$ の場合、最大 $B_i$ 個まで、そこの'1'を移動元として確定させることができる。 これは、右方向に累積和(ただし $B_i$ 回まで)を取る操作となる。

i=1           (B1を移動元とした場合のみ考慮)
j\k  0  1  2
0    1  -  -
1    1  1  -
2    1  1  1

$DP[i]$ は、$DP[i-1]$ の下方向の累積和と、右方向の累積和を合計したものとなる。

i=1              C0 C1  (j, k)   C0 C1  (j, k)
j\k  0  1  2
0    1  -  -      0  2  (0, 0)    1  1  (1, 1)
1    2  1  -      0  3  (1, 0)    1  2  (1, 0)
2    3  1  1      0  4  (2, 0)    1  3  (2, 0)
                  2  0  (2, 2)
                  2  1  (2, 1)
                  2  2  (2, 0)

「ある $B_i$ を移動元とし、さらに移動先とする」という操作は無駄なので、 $j$ を増やすか $k$ を増やすか、どちらかのみ考慮すればよい。

最終的に、移動元・先の個数の辻褄の合っている、対角線上の数の合計が、答えとなる。

from itertools import groupby
import numpy as np


def solve(s, k):
    sgrp = [len(list(itr)) for c, itr in groupby(s)]
    if len(sgrp) == 1:
        return 1

    one_bucket = []
    MOD = 998244353

    if s[0] == '1':
        one_bucket.append(sgrp[0])
        sgrp = sgrp[1:]
    else:
        one_bucket.append(0)

    cur = 0
    for l in sgrp:
        if cur == 0:
            one_bucket.extend([0] * (l - 1))
        else:
            one_bucket.append(l)
        cur ^= 1

    k = min(k, sum(one_bucket))
    dp = np.zeros((k + 1, k + 1), dtype=np.int64)
    dp[0, 0] = 1

    for c in one_bucket:
        ndp = np.add.accumulate(dp, axis=0)
        if c > 0:
            for j in range(1, c + 1):
                ndp[:, j:] += dp[:, :-j]
            ndp = np.tril(ndp)
        dp = ndp % MOD

    return np.diag(dp).sum() % MOD


s, k = input().split()
k = int(k)

print(solve(s, k))

D - Secret Passage

問題

  • '0' と '1' のみからなる文字列 $S$ が与えられる
  • 以下の操作を0回以上、任意の回数繰り返す
  • 操作
    • $S$ の先頭2文字を取り除き、そのうち片方を捨て、もう片方を $S$ の任意の位置に挿入する
    • $S$ が1文字の時は行えない
  • 最終的にできる可能性のある文字列の個数を 998244353 で割った余りを求めよ
  • $1 \le |S| \le 300$

解法

2種類の3次元DP。

最終的な文字列を $T$ とすると、これを構成する文字は「操作によって挿入されたもの」「一度も操作を受けなかったもの」の2つである。

先頭からどこまでを一度は操作対象として取り除いたか、を固定する。 先頭 $i$ 文字を取り除いたとして、$S$ をそこで分けて $P_i,Q_i$ とする。

「$P_i$ から挿入用に残せる $(0,1)$ の個数の組」それぞれに対し、「それを $Q_i$ に挿入したときの場合の数」を足し合わせれば、$i$ を固定したときの答えが求まる。

言い換えると、

  • $DP1[i][j][k]=P_i$ から $(0,1)$ を $(j,k)$ 個残せるか(bool)
  • $DP2[i][j][k]=Q_i$ に $(0,1)$ を $(j,k)$ 個挿入してできる文字列の数

として、$DP1[i][j][k]=true$ であるような $(j,k)$ につき $DP2[i][j][k]$ を合計すると、$i$ を固定したときの答えが求まる。

$i$ が同じで $(j,k)$ が異なるものは、最終的な $T$ を構成する $(0,1)$ の個数が全て異なるので、異なる $(j,k)$ 間で同じ $T$ を重複して数えてしまうことはない。

DP1

結構いろんな遷移パターンがあり、漏らさぬよう注意深く考える必要がある。

$i$ が小さい方から考えて、たとえば $DP1[i][j][k] = true$ であるとわかっているとする。$DP1[i+1],DP1[i+2]$ への遷移を考える。

■S[i+1] or S[i+2] に'0'が含まれる場合
  DP1[i] で残す j個、k個はそのまま、新たに次の2文字に操作して、'0' の方を残せる
    DP1[i+2][j+1][k] = true

■S[i+1] or S[i+2] に'1'が含まれる場合
  DP1[i] で残す j個、k個はそのまま、新たに次の2文字に操作して、'1' の方を残せる
    DP1[i+2][j][k+1] = true

■j > 0 or k > 0 の場合
  残すことにしている文字の1つ c を i 番目に挿入しておき、c と S[i+1] に対して操作、c を残す
  残す文字は変えず、S[i+1] を消費できる
    DP1[i+1][j][k] = true

■S[i+1] = '0'、k > 0 の場合
  残すことにしている '1' の1つを i 番目に挿入しておき、'1' と S[i+1] に対して操作、S[i+1] を残す
  '1' を1つ犠牲に、'0' を1つ多く残せる
    DP1[i+1][j+1][k-1] = true

■S[i+1] = '1'、j > 0 の場合
  ↑と逆
    DP1[i+1][j-1][k+1] = true

DP2

異なる挿入の仕方でも同じ文字列が出来上がる場合があり、それらを重複して数えないよう注意する。

これは、問題Cと同じで「最終的な $T$ のうち、$Q$ は左から貪欲に出現するとし、残りを挿入したと考える」と、上手く重複を除ける。

Q: 101    挿入する文字: {0,0,1,1}

Tの一例: 0100111
          ^^ ^    (^: Q 由来と見なす文字 )

これを念頭に置きつつ、$i$ が大きい方から埋めていく。

$DP2[i+1][j][k]$ からの、$DP2[i]$ への遷移を考える。

$Q_i$ は、$S$ の $i+1$ 文字目を $Q_{i+1}$ の先頭に加えた文字列となる。ここでは '0' だったとする。

        i
S ......00101
Q[i+1]    101
Q[i]     0101

$DP2[i][j][k]$ によって数え挙げられる文字列の1つを $T_{i,j,k}$ とする。

「$T$ の中で、$Q$ は左から貪欲に出現すると見なす」ルールがあるので、個々の $T_{i+1,j,k}$ の中で $Q_{i+1}$ が出現する(と見なす)位置は一意である。

Q[i+1]:  101    j=2  k=2

T: 0100111    1100011    0011101
    ^^ ^      ^ ^  ^       ^  ^^

$DP2[i+1][j'][k']$ から $DP2[i][j][k]$ への遷移を考える時は、

  • $T_{i,j,k}$ の末尾に $T_{i+1,j',k'}$ が出現する
  • $T_{i,j,k}$ の中で $Q_{i}$ が出現する位置を見ると、$Q_{i+1}$ の部分の出現位置は $T_{i+1,j',k'}$ と同じ

もののみを数える。

$T_{i+1,j,k}$ の先頭に '0' を付けたものは、$T_{i,j,k}$ として条件を満たす。

Q[i]:  0101     j=2  k=2

T: 00100111    01100011    00011101
   ^ ^^ ^      ^^ ^  ^     ^  ^  ^^

また、その前に '1' を複数個追加してもよい。

Q[i]:  0101     j=2  k=5(3個追加)

T: 11100100111    11101100011    11100011101
      ^ ^^ ^         ^^ ^  ^        ^  ^  ^^

あと、一応 '0' の後ろにさらなる '0' を複数個追加するのもOKなのだが……

Q[i]:  0101     j=5  k=2

T: 00000100111    00001100011    00000011101
   ^    ^^ ^      ^   ^ ^  ^     ^     ^  ^^

これは、$DP2[i+2]$ から $DP2[i+1]$ に遷移する時に、1つ上のように「'0' を複数個追加してもよい」として既に考慮されているので、改めて数えなくてよい。

そして、これ以外の場所に '0' や '1' を追加すると $Q_{i+1}$ が出現する位置が崩れてしまう。

Q[i]:  0101     j=3  k=5

        _______  T[i+1,2,2] と一致する部分
T: 110010100111
     ^ ^^^       この T の中での Q[i] の出現位置 
    
T[i+1,2,2] における '101' の出現位置と異なるため DP[i+1][2][2] からの遷移では数え挙げない


この T は、DP2[i+1][3][3] において 010100111 が含まれるので、ここから派生するものとして数えられる

T[i+1,3,3]:    010100111
                ^^^
T[i,3,5]:   110010100111
              ^ ^^^

まとめると、$Q_i$ の先頭文字が '0' なら、$T_{i+1,j',k'}$ の先頭に0個以上の '1' と1個の '0' を追加したものが、$DP2[i+1][j'][k']$ から遷移できる $DP2[i][j][k]$ となる。

DP上の計算で言うと、$DP2[i][j][k],DP2[i][j][k+1],DP2[i][j][k+2],...$ のそれぞれに、$DP2[i+1][j][k]$ が加算される。

もらうDPに変換すると、$\displaystyle DP2[i][j][k] = \sum_{k'=0}^{k}DP2[i+1][j][k']$ となる。

$Q_i$ の先頭文字が '1' なら逆に $\displaystyle DP2[i][j][k] = \sum_{j'=0}^{j}DP2[i+1][j'][k]$ となる。

異なる $i$ 間での重複除外

以上で同一の $i$ については考察ができたが、$i$ が異なる場合でも結果的に同じ文字列になるものがある。この重複を除く必要がある。

たとえば、以下の2つがあったとして

  • (A) $DP2[i+2][2][1]$ … 011 に 0,0,1 を挿入する
  • (B) $DP2[i][1][0]$ … 10011 に 0 を挿入する

どちらも、010011、100101などが重複してできる。

この時、挿入する文字が多い(A)の方が自由度が高く、(B)から作れる文字列は全て(A)でも作れる。よって、(A) だけを数え挙げるようにすれば、重複は除外できる。

$i$ が大きい方から計算し、ある $(i,j,k)$ の組でできる $T$ の $(0,1)$ の個数を $(a,b)$ とすると、 より小さい $i$ において、そこからできる $T$ に含まれる $(0,1)$ の個数が $(a,b)$ のものは除外するとよい。

まとめ

  • $DP1$ を計算
  • $DP2$ を計算
  • $DP1$ を $i$ の大きい方から見て、$DP1[i][j][k]=true$ なら
    • $DP2[i][j][k]$ を答えに加算
    • $(0,1)$ の個数が同じになる、より $i$ の小さい $DP1[i][j][k]$ を false にする

以上で答えが求まる。

import numpy as np
import numba as nb


@nb.njit('(i1[:],)', cache=True)
def solve(s):
    n = len(s)
    size = n // 2
    MOD = 998244353
    dp_can = np.zeros((n + 1, size + 1, size + 1), dtype=np.int8)
    dp_cnt = np.zeros((n + 1, size + 1, size + 1), dtype=np.int64)

    dp_cnt[n, 0] = 1
    for j in range(size):
        dp_cnt[n, j + 1] = np.cumsum(dp_cnt[n, j]) % MOD

    for i in range(n - 1, -1, -1):
        if s[i] == 0:
            for j in range(size + 1):
                dp_cnt[i, j] = np.cumsum(dp_cnt[i + 1, j]) % MOD
        else:
            for k in range(size + 1):
                dp_cnt[i, :, k] = np.cumsum(dp_cnt[i + 1, :, k]) % MOD

    dp_can[0, 0, 0] = 1
    for i in range(n):
        dp_can[i + 1] |= dp_can[i]
        if s[i] == 0:
            dp_can[i + 1, 1:, :-1] |= dp_can[i, :-1, 1:]
        else:
            dp_can[i + 1, :-1, 1:] |= dp_can[i, 1:, :-1]
        if i < n - 1:
            if s[i] == 0 or s[i + 1] == 0:
                dp_can[i + 2, 1:, :] |= dp_can[i, :-1, :]
            if s[i] == 1 or s[i + 1] == 1:
                dp_can[i + 2, :, 1:] |= dp_can[i, :, :-1]
    dp_can[1:, 0, 0] = 0

    ans = 0

    for i in range(n, -1, -1):
        for j in range(size + 1):
            for k in range(size + 1):
                if dp_can[i, j, k] == 0:
                    continue

                ans = (ans + dp_cnt[i, j, k]) % MOD

                a, b = j, k
                for h in range(i - 1, -1, -1):
                    if s[h] == 0:
                        a -= 1
                    else:
                        b -= 1
                    if a < 0 or b < 0:
                        break
                    dp_can[h, a, b] = 0

    return ans


s = np.array([int(c) for c in input()], dtype=np.int8)
print(solve(s))

E - Permutation Cover

問題

  • 整数 $K$ と整数列 $A_1,A_2,...,A_K$ が与えられる
  • 以下の数列 $P$ を構築せよ
    • 全ての項は $1~K$ の整数からなる
    • $i=1~K$ について、$i$ をちょうど $A_i$ 個含む
    • 全ての項について、その項を含む長さ $K$ の連続する部分列で $1~K$ の順列になっているものが存在する
  • 構築可能な場合、その中で辞書順最小のものを出力せよ
  • 不可能な場合、-1を出力せよ
  • $1 \le K \le 100$
  • $(A_i の総和) \le 1000$

$K = 3, A = [2, 4, 3]$

2 3 1 2 3 2 3 1 2

どの項を見ても、どれかしらの (1, 2, 3) の順列に含まれている

2 3 1
    1 2 3
          2 3 1
            3 1 2

ただしこれは辞書順最小では無い。

解法

可能不可能はすぐわかるが、辞書順最小の達成が全然分からなかった。

可能判定

なんとなく、$A_i$ 間の差が激しすぎたら無理そう。

もう少し具体的に考える。'2' をなるべくたくさん使い、'1' をなるべく少なく使おうとすると、

2 _ _ 1 _ _ 2 2 _ _ 1 _ _ 2
`-----'-----' `-----'-----'
  順列  順列    順列  順列

みたいに'2'を'1'で挟むと、2倍までは使える。 2倍を超えると、余った'2'から、他の'2'を介さずに'1'を含められなくなるので不可能。

格差が2倍未満や他の個数が入り交じっている場合なら、

A B C D E F A B C   A B C D E F A B
===== ----- =====   === ------- ===

みたいにして、多く消費したいやつ(=)で、 なるべく消費したくないやつ(-)をサンドすることで 1ブロックあたり2個消費するか1個消費するかを自由に選べるので、必ず構築できる。

なので、「最大値が、最小値の2倍を超えない」が条件となる。

辞書順最小の構築方法

誤った方法

判定が上記のようにできてしまうもんだから、 つい「ブロックを作って、それを繋げる」方向で考えてしまいたくなるが、 それでは上手くいかない。

A=[4, 7, 5]

ブロック分けでの最善
1 2 3|2 1 3 2|2 1 3 2|2 3 1 2 3

真の最善
1 2 3 2 1 2 3 2 1 2 3 2 3 1 2 3

真の最善は、ブロックに分割しようとしても上手くできない。

1 2 3       が確定した後、前のブロックの末尾を利用して
    3 2 1   というブロックを作っている、とも言える?

ブロックごと独立の構築では、前のブロックの末尾の項を利用する、 ということができないため正しくならない。

解説の方法

まず、順列になっている部分列(の末尾)は、$K$ 個以内おきに必ず出現する。 (そのような部分列で全体が覆われてないといけないので)

    o o     o o o  ←順列になっている部分列の末尾
2 1 3 2 2 3 1 2 3
~~~~~   ~~~~~
  ~~~~~   ~~~~~
            ~~~~~

そこで「末尾が常に順列」となる状態を保ちつつ、暫定の数列に $1~K$ 文字追加することを繰り返して順次確定させていく。

2 1 3
2 1 3 2
2 1 3 2 2 3 1
2 1 3 2 2 3 1 2
2 1 3 2 2 3 1 2 3

何文字新たに付け加えるか(暫定配列の末尾何文字をダブらせるか)を $1~K$ 個の全通り試す。

そのうちどれを採用するかは、 「それを確定させた後、さらに残りを破綻せずに配置できる」もののうち、辞書順最小のものを使う。

暫定配列の末尾が順列であり、次も順列を保つように遷移するので「$n$ 文字追加するときは、使う文字の集合はこれこれ」というのが一意に決まる。

たとえば、サンプル2において「1 2 3 4 1」まで決まっているとして、次の遷移で残りを破綻しないように拡張するには、以下のようになる。 破綻しない条件は後述。

A = [3, 2, 3, 2]

(1文字)  1 2 3 4 1 [x]        置くなら2だが、破綻
(2文字)  1 2 3 4 1 [3 2]
(3文字)  1 2 3 4 1 [3 2 4]
(4文字)  1 2 3 4 1 [3 1 2 4]

従って、次はこの中で辞書順最小である「3 1 2 4」を拡張する。

辞書順最小は、このように 「先頭から考えて、残りで破綻せずに条件を満たせることがわかりさえすれば、その中で一番いい候補を確定してOK」 という方法が使える。 (今回の場合、それを「順列となる部分列」ごとに確定させる、というのがポイント)

後続が破綻しない条件

さて、では、どういう条件なら残りを破綻せずに配置できるのか。 これは、可能判定の議論と同様、

  • 置いた後の $A_i$ の状態が、だいたい $最大値 \le 2 \times 最小値$

となっていればよいが、違うのは、確定済みの項が存在するということ。

末尾が上手く配置されていれば(そのように確定していけば)、もう少し条件が緩くなる。

たとえば $K=2$ として'2'を'1'の約2倍消費しなければならないとして、

確定済 | 考慮中  | その後
 _ _ _ | 2 1     | 2  2 1 2  2 1 2 ...

このように、【考慮中】で '1' を後の方に置くようにすれば、 【その後】において '2'を '1' の2倍よりもう1つ多く置くことができる。

このように置かなければいけないのは、ギリギリの時、 つまり置いた後の $A_i$ の状態が、$最大値 = 2 \times 最小値 + 1$ となる時。

どのように置けばよいかというと、その時点で $A_i$ が最大値をとる全ての $i$ を、最小値をとる全ての $i$ の前に置くとよい。

また、ギリギリで無い場合、

  • $最大値 \gt 2 \times 最小値 + 1$ ならその拡張は不可能
  • $最大値 \lt 2 \times 最小値 + 1$ なら自由に並べてよい、つまり昇順が最適

計算量

最終的な数列長($A_i$ の総和)を $T$とする。

この貪欲法は1文字ずつしか拡張されない場合もあるので、「$1~K$ 文字拡張できるか試す」という操作は $O(T)$ 回繰り返されうる。

$k$ 文字拡張できるか調べる操作では $K$ 要素のイテレートと $k$ 要素のソートを行う1)ので、$O(K+k \log{k})$ で、これを $k=1~K$ まで繰り返す。

従って、全体でだいたい $O(TK^2)$ かかる。

import sys


def find_permutation(aaa, use):

    max_a = -1
    min_a = 1005
    max_fixed = -1

    for i in range(k):
        a = aaa[i]
        if i + 1 in use:
            min_a = min(min_a, a)
            max_a = max(max_a, a)
        else:
            max_fixed = max(max_fixed, a)

    if max(max_a, max_fixed + 1) > 2 * min_a:
        return None

    if max_a < 2 * min_a:
        return sorted(use)

    front = []
    rear = []
    either = []
    for i in use:
        if aaa[i - 1] == max_a:
            front.append(i)
        elif aaa[i - 1] == min_a:
            rear.append(i)
        else:
            either.append(i)

    max_front = front[-1]
    for i in either:
        if i < max_front:
            front.append(i)
        else:
            rear.append(i)
    front.sort()
    rear.sort()
    front.extend(rear)

    return front


def solve(k, aaa):
    if k == 1:
        return [1] * aaa[0]

    min_a = min(aaa)
    max_a = max(aaa)
    if min_a * 2 < max_a:
        return [-1]

    ans = []
    ans.extend(find_permutation(aaa, set(range(1, k + 1))))
    for i in range(k):
        aaa[i] -= 1

    remaining = sum(aaa)
    while remaining:
        use = set(range(1, k + 1))
        candidates = []
        for r in range(k):
            result = find_permutation(aaa, use)
            if result is not None:
                candidates.append(result)

            use.remove(ans[-r - 1])
        adopted = min(candidates)
        ans.extend(adopted)
        for i in adopted:
            aaa[i - 1] -= 1
        remaining -= len(adopted)

    return ans


k, *aaa = map(int, sys.stdin.buffer.read().split())
print(*solve(k, aaa))

1)
データの持ち方次第で必須ではない
programming_algorithm/contest_history/atcoder/2020/0620_agc046.txt · 最終更新: 2020/08/21 by ikatakos
CC Attribution 4.0 International
Driven by DokuWiki Recent changes RSS feed Valid CSS Valid XHTML 1.0