純粋培養プログラマが就活を意識して Python で ABC を解く [Day 1]

このシリーズの目的

就活を少しだけ意識して綺麗な Python コードを書けるようになる。そのため、出来るだけ綺麗で分かりやすいコードを書くことが目的。（そのため、for文でゴリ押しして PyPy で提出、といったことはしない。）

今回解くコンテスト

ABC 203 (Sponsered by Panasonic)

Tasks - AtCoder Beginner Contest 203（Sponsored by Panasonic）

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

A問題

同じ値があるかないかどうかは、ソートして調べると簡単に分かる。

// 解答
a, b, c = sorted(map(int, read().split()))
 
if a == b:
    print(c)
elif b == c:
    print(a)
else:
    print(0)

// 解答

a, b, c = sorted(map(int, read().split()))

if a == b:

print(c)

elif b == c:

print(a)

else:

print(0)

B問題

やるだけ。

// 解答
N, K = map(int, input().split())
 
sum = 0
for i in range(N):
    for j in range(K):
        sum += i * 100 + j + 101
print(sum)

// 解答

N, K = map(int, input().split())

sum = 0

for i in range(N):

for j in range(K):

sum += i * 100 + j + 101

print(sum)

C問題

ペアに対して、「第１要素を優先したソーティング」の処理が要求される。

// 解答
N, K = map(int, input().split())
 
friends = []
for i in range(N):
    a, b = map(int, input().split())
    friends.append((a, b))
 
friends.sort()
 
pos = 0
for (x, given) in friends:
    if x - pos <= K:
        K -= (x - pos)
        pos = x
        K += given
    else:
        break
 
print(pos + K)

// 解答

N, K = map(int, input().split())

friends = []

for i in range(N):

a, b = map(int, input().split())

friends.append((a, b))

friends.sort()

pos = 0

for (x, given) in friends:

if x - pos <= K:

K -= (x - pos)

pos = x

K += given

else:

break

print(pos + K)

D問題

苦戦した。二分探索をした後は、累積和を使って高速に数を数えるだけなのだが、その処理を純粋に実装するとTLEとなる。理由としては、Pythonの二重ループが遅いためであると考えられる。

// TLE回答
import numpy as np
N, K = map(int, input().split())
A = np.array([list(map(int, input().split())) for i in range(N)])
ok, ng = 10 ** 9 + 1, -1
while ok - ng > 1:
    # 中央値を mi 以下にできるか？
    mi = (ok + ng) // 2
    B = np.where(A.reshape(N ** 2) <= mi, 1, 0).reshape((N, N))
    C = np.zeros(shape=(N+1,N+1))

    for i in range(N):
        for j in range(N):
            C[i + 1][j + 1] += B[i][j]
    
    for i in range(N):
        for j in range(N-1):
            C[i + 1][j + 2] += C[i + 1][j + 1]
 
    for i in range(N-1):
        for j in range(N):
            C[i + 2][j + 1] += C[i + 1][j + 1]
    
    flag = False
    for i in range(K, N+1):
        for j in range(K, N+1):
            # (i, j) を右下にする四角が条件を満たすか？
            cnt = C[i][j] - C[i - K][j] - C[i][j - K] + C[i - K][j - K]
            if cnt >= (K ** 2 + 1) // 2:
                flag = True

    if flag:
        ok = mi
    else:
        ng = mi
 
print(ok)

// TLE回答

import numpy as np

N, K = map(int, input().split())

A = np.array([list(map(int, input().split())) for i in range(N)])

ok, ng = 10 ** 9 + 1, -1

while ok - ng > 1:

# 中央値を mi 以下にできるか？

mi = (ok + ng) // 2

B = np.where(A.reshape(N ** 2) <= mi, 1, 0).reshape((N, N))

C = np.zeros(shape=(N+1,N+1))

for i in range(N):

for j in range(N):

C[i + 1][j + 1] += B[i][j]

for i in range(N):

for j in range(N-1):

C[i + 1][j + 2] += C[i + 1][j + 1]

for i in range(N-1):

for j in range(N):

C[i + 2][j + 1] += C[i + 1][j + 1]

flag = False

for i in range(K, N+1):

for j in range(K, N+1):

# (i, j) を右下にする四角が条件を満たすか？

cnt = C[i][j] - C[i - K][j] - C[i][j - K] + C[i - K][j - K]

if cnt >= (K ** 2 + 1) // 2:

flag = True

if flag:

ok = mi

else:

ng = mi

print(ok)

ACするために今回行った工夫は、以下の通り。全体として、「ループを素直に実装しない」という方針で工夫を行なった。

配列の確保の仕方
- C++なら適当な配列を用意するために普通にループを回していたが、今回はnp.where （〇〇な部分は■■にする、といった処理）と np.insert（0-indexの配列を1-indexに移すために使用）を用いた。
累積和の計算
- ループを回して累積和を計算するのではなく np.cumsum を用いる
スライスを用いた高速な計算
- 各 \(K*K\) 領域に含まれる数の総和がいくつかを高速に判定したい。そのために、切り出した二次元配列同士の加算・減算に帰着させた。

// AC回答（やってることは同じ）
import numpy as np
N, K = map(int, input().split())
A = np.array([list(map(int, input().split())) for i in range(N)])
 
ok, ng = 10 ** 9, -1
 
while ok - ng > 1:
    # 中央値を mi 以下にできるか？
    mi = (ok + ng) // 2
    
    B = np.where(A <= mi, 1, 0)
    
    C = np.insert(B, 0, 0, axis=0)
    C = np.insert(C, 0, 0, axis=1)
    C = np.cumsum(C, axis=0)
    C = np.cumsum(C, axis=1)
    res = C[K:N+1, K:N+1] - C[0:N-K+1, K:N+1] - C[K:N+1, 0:N-K+1] + C[0:N-K+1, 0:N-K+1]
    D = np.where(res >= (K ** 2 + 1) // 2, 1, 0)
    if D.sum() >= 1:
        ok = mi
    else:
        ng = mi
 
print(ok)

// AC回答（やってることは同じ）

import numpy as np

N, K = map(int, input().split())

A = np.array([list(map(int, input().split())) for i in range(N)])

ok, ng = 10 ** 9, -1

while ok - ng > 1:

# 中央値を mi 以下にできるか？

mi = (ok + ng) // 2

B = np.where(A <= mi, 1, 0)

C = np.insert(B, 0, 0, axis=0)

C = np.insert(C, 0, 0, axis=1)

C = np.cumsum(C, axis=0)

C = np.cumsum(C, axis=1)

res = C[K:N+1, K:N+1] - C[0:N-K+1, K:N+1] - C[K:N+1, 0:N-K+1] + C[0:N-K+1, 0:N-K+1]

D = np.where(res >= (K ** 2 + 1) // 2, 1, 0)

if D.sum() >= 1:

ok = mi

else:

ng = mi

print(ok)

感想

numpyやリスト内包表記にもっと詳しくなりたいです。（普段Pythonで競プロしてる人ってこんな大変なことやってるんですか……？）