本記事は、AtCoder Beginner Contest 457 (Polaris.AIプログラミングコンテスト) に参加した際の、A〜D問題の復習と解答の備忘録です。コンテスト中に考えた解法の方針や、提出したPythonのコードについて整理しています。

A - Array /

実行時間制限: 2 sec / メモリ制限: 1024 MiB

配点 :
100 点

問題文

長さ $N$ の数列 $A=(A_1, A_2, \dots, A_N)$ が与えられます。

その後に 1 以上 $N$ 以下の整数 $X$ が与えられます。

$A_X$ の値を出力してください。

制約

$1 \le X \le N \le 100$
$1 \le A_i \le 100$
入力される値は全て整数

自分の解答の方針

$A$を配列として受け取り、添え字が$X-1$のものを出力

提出したコード

N = int(input())
A = list(map(int,input().split()))
X = int(input())
print(A[X-1])

B - Arrays /

実行時間制限: 2 sec / メモリ制限: 1024 MiB

配点 :
200 点

問題文

$N$ 個の数列 $A_1, A_2, \dots, A_N$ が与えられます。

数列 $A_i$ の長さは $L_i$ であり $A_i = (A_{i,1}, A_{i,2}, \dots, A_{i,L_i})$ です。

その後、整数 $X,Y$ が与えられます。 $A_{X,Y}$ の値を出力してください。

制約

$1 \le N \le 2 \times 10^5$
$1 \le L_i$
$L_i$ の総和は $2 \times 10^5$ 以下
$1 \le A_{i,j} \le 1000$
$1 \le X \le N$
$1 \le Y \le L_X$
入力される値はすべて整数

自分の解答の方針

$A$を$L$を含めて二次元配列として受け取り、添え字が$X-1,Y$のものを出力

提出したコード

N = int(input())
A = [list(map(int,input().split())) for _ in range(N)]
X,Y = map(int,input().split())
print(A[X-1][Y])

C - Long Sequence /

実行時間制限: 2 sec / メモリ制限: 1024 MiB

配点 :
300 点

問題文

整数 $N,K$ と $N$ 個の整数列 $A_1, A_2, \dots, A_N$ 、長さ $N$ の整数列 $C=(C_1, C_2, \dots, C_N)$ が与えられます。整数列 $A_i$ の長さは $L_i$ であり $A_i = (A_{i,1}, A_{i,2}, \dots, A_{i,L_i})$ です。また、 $1 \le K \le \sum_{i=1}^{N} C_i L_i$ が保証されます。

$A$ と $C$ から以下の手順で整数列 $B=(B_1, B_2, \dots, B_{\sum_{i=1}^{N} C_i L_i})$ を構成します。

$B$ を長さ 0 の整数列とする。
$i=1,2,\dots,N$ の順に以下の操作を行う：
- $B$ の末尾に $A_i$ を連結させる操作を $C_i$ 回行う。

$B_K$ の値を求めてください。

制約

$1 \le N$
$1 \le L_i$
$\sum_{i=1}^{N} L_i \le 2 \times 10^5$
$1 \le A_{i,j} \le 10^9$ $(1 \le j \le L_i)$
$1 \le C_i \le 10^9$
$1 \le K \le \sum_{i=1}^{N} C_i L_i$
入力される値は全て整数

１〜３回目の提出で考えたこと

問題にシグマが出ていてビビった
とりあえず入出力例から見て、$K$が大きければ$K$を$C_iL_i$ずつ減らしていき、$C_iL_i$が$K$より大きければ$K$を$L_i$で割ったあまりから添え字を計算して出力すればいいと思った。
二個目のテストケースのように余りが0になるときに$A[i][0]$となり、$L_i$を参照して間違いになるので、無理やり条件分岐で$A[i][1]$にしたりして通ったものを提出した。

提出したコード

N,K = map(int,input().split())
A = [list(map(int,input().split())) for _ in range(N)]
C = list(map(int,input().split()))
for i in range(N):
    if C[i]*A[i][0] < K:
        K -= C[i]*A[i][0]
    else:
        if K % A[i][0] == 0:
            print(A[i][1])
        else:
            print(A[i][K%A[i][0]])
        break

結果:WA

４回目の提出で考えたこと

一旦落ち着いて添え字について考えることにした。
$K$の添え字がややこしいので、初めに-1をして、0-indexedに揃えた。
そして、あまりに1を足すことで、$L_i$を参照することがなくなった。

提出したコード

N,K = map(int,input().split())
A = [list(map(int,input().split())) for _ in range(N)]
C = list(map(int,input().split()))
K -= 1
for i in range(N):
    if C[i]*A[i][0] <= K:
        K -= C[i]*A[i][0]
    else:
        print(A[i][K%A[i][0]+1])
        break

結果:AC

D - Raise Minimum /

実行時間制限: 2 sec / メモリ制限: 1024 MiB

配点 :
400 点

問題文

長さ $N$ の数列 $A=(A_1, A_2, \dots, A_N)$ と整数 $K$ が与えられます。

あなたは次の操作を 0 回以上 $K$ 回以下行うことができます。

$1 \le i \le N$ を満たす整数 $i$ を 1 つ選び、 $A_i$ に $i$ を加える。

操作後の数列に対する $\min_{1 \le i \le N} A_i$ としてありうる最大値を求めてください。

制約

$1 \le N \le 2 \times 10^5$
$1 \le A_i \le 10^{18}$
$1 \le K \le 10^{18}$
入力される値はすべて整数

一度目の提出で考えたこと

C問題で誤答を含めて時間をとられすぎてしまったことに焦り、制約を見逃してしまった。
heapqを使って値と添え字を保持し、最小値を取り出して添え字を足して戻す操作を$K$回繰り返せばいいと思った。

提出したコード

import heapq
N,K = map(int,input().split())
A = list(map(int,input().split()))
B = [(A[i],i+1) for i in range(N)]
heapq.heapify(B)
for i in range(K):
    C,idx = heapq.heappop(B)
    heapq.heappush(B,(C+idx,idx))
ans,_ = heapq.heappop(B)
print(ans)

結果:TLE

２度目の提出で考えたこと

問題と制約を読み直したところ、最小値の最大化なので答えで二分探索する方法が使えそうだと思った。操作後の数列の全ての要素を$K$回以内の操作である値にすることができるかどうかを判定すればいい。各$A_i$について、目標値が定まっていれば割り算１回で操作の回数を計算できる。この時、あまりを切り上げることに注意したほうがよさそう。合計で計算量も$NlogK$くらいになりそうで、大丈夫そう。二分探索の範囲は、$A_i$が$1 \le A_i \le 10^{18}$なので、それに合わせる

提出したコード

def check(num):
    sum = 0
    for i in range(len(A)):
        if A[i] < num:
            sum += (num-A[i]+i)//(i+1)
    if sum <= K:
        return True
    else:
        return False


N,K = map(int,input().split())
A = list(map(int,input().split()))
ok = 1
ng = 10**18
while (ng-ok>1):
    middle = (ok+ng)//2
    if check(middle):
        ok = middle
    else:
        ng = middle
print(ok)

結果:WA

３度目の提出で考えたこと

ロジック自体はあってそう
制約を見直したら、答えは$10^{18}$以上になることがありそうなので、これが原因っぽい
$10^{18}$に1を$10^{18}$回足しても$2*10^{18}$にしかならないので、ちょっと余裕をもって$10^{20}$くらいにしておけば大丈夫そう

提出したコード

def check(num):
    sum = 0
    for i in range(len(A)):
        if A[i] < num:
            sum += (num-A[i]+i)//(i+1)
    if sum <= K:
        return True
    else:
        return False


N,K = map(int,input().split())
A = list(map(int,input().split()))
ok = 1
ng = 10**20
while (ng-ok>1):
    middle = (ok+ng)//2
    if check(middle):
        ok = middle
    else:
        ng = middle
print(ok)

結果:AC