paizaプログラミング問題Rank S (Python)

paiza×Qiita記事投稿キャンペーンに参加しました．
関連リポジトリは，以下になります．

ここでは，Rank S問題の解答と簡単な解説を公開します．
言語はPythonを選択しました．結果は以下のようになりました．

Rank S

村人の友好関係

group_popularity.py

import numpy as np

N, M, Q = map(int, input().split())
N += 1  # 入力が1～Nのため

m = np.zeros((N, N), dtype=int)
for _ in range(M):
    a, b, f = map(int, input().split())
    m[a][b] = m[b][a] = f

g = np.zeros(N, dtype=bool)
for _ in range(Q):
    q = input().split()
    g[int(q[1])] = q[0] == "+"
    print(np.max(m[g][:, ~g], initial=0))

pythonの配列が0始まりであるのに対して，入力が1～Nのため，相性が悪いです．
入力が来るたびに-1するのを避けるため，Nを+1して配列作成に利用します．

mは，村人間の友好度を格納した2次元配列です．
村人aとbの友好度f = m[a][b] = m[b][a]です．

gは，同好会グループに入会している村人を表す1次元配列です．
g[i]がTrueの場合，村人iは同好会グループに入会しています．

m[g][:, ~g]は，グループ内の村人とグループ外の村人の間の友好度です．
最大値の算出や配列の反転は，numpyの機能を利用しています．

島探し

search-island.py

from scipy.ndimage import label

f = lambda: list(map(int, input().split()))
m = [f() for _ in range(f()[1])]
print(label(m)[1])

scipy.ndimageのlabel関数を使います．
分かり易さのため，標準入力1行を数値リストとして受け取る処理を関数化しています．
2次元配列mをlabel関数に渡すと，欲しい出力が2番目に出てきます．

mod7占い

mod7.py

from itertools import combinations

a = [int(input()) for _ in range(int(input()))]
print(sum(not sum(x) % 7 for x in combinations(a, 3)))

itertoolsのcombinationsを利用します．
まず，入力を数値リストaとして受け取ります．
そして，リストaから取り出した3要素について，その合計が7の倍数であればTrue，
7の倍数でなければFalseとします．
リストaから3要素を取り出す全ての組み合わせを試し，Trueの合計を出力します．

文字列収集

word_collection.py

from collections import defaultdict

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

d = defaultdict(int)
for _ in range(N):
    k, v = input().split()
    d[k] += int(v)

for _ in range(M):
    q = input()
    print(sum(v * k.startswith(q) for k, v in d.items()))

collectionsのdefaultdictを利用します．
文字列とその価格は，辞書型のデータdとして保存します．
同じ文字列が複数売っていた場合，その文字列の価格は，合計して1つにまとめます．
例えば，abcd 14とabcd 6が入力にあった場合，d["abcd"] ⇒ 20とします．

各クエリ文字列qについて，dの全データと先頭の文字列が一致するか検証します．
itemsを使うと，辞書のkeyとvalueを取り出してループできます．
自分は掛け算の方が好みですが，以下のようにif文を書く方法もあります．

    print(sum(v for k, v in d.items() if k.startswith(q)))

十億連勝

continuous_winning.py

from collections import Counter

N, X = map(int, input().split())
A = [int(input()) for _ in range(N)]

# b：(現在の連勝数, X連勝後に敗北) とその数
b = Counter({(0, False): 1})

for a in A:  # 各ステージの処理
    c = Counter()
    for (w, t), n in b.items():
        if w + a <= X:
            # 全勝しても連勝数がX以下の場合
            c[(w + a, t)] += n  # 全勝パターン
            c[(0, t)] += n * a  # 途中敗北
        else:
            # 全勝すると連勝数がXを超える場合
            c[(0, True)] += n  # X連勝後に敗北
            c[(0, t)] += n * (X - w)  # X連勝以下で敗北
            # n * (w + a - X) パターン # 連勝数がX以上

    b = c  # 状態を更新

print(sum(n for (w, t), n in b.items() if w == X or t) % 10**9)

collectionsのCounterを利用します．
bは，(現在の連勝数, X連勝後に敗北) とその数が保存されています．
ステージ単位でbを更新し，最後に，X連勝中 or X連勝後に敗北のパターン数の合計を出力します．出力は下9桁に制限します．

ステージaでは，まず，空のデータcを用意します．
そして，1つ前のステージでの状態別に，ステージa終了後のパターンを計算します．
全勝しても連勝数がX未満の場合，考えるべきなのは，全勝パターンと途中敗北の2つです．
全勝すると連勝数がXを超える場合，考えるべきなのは， X連勝後に敗北パターンとX連勝未満で敗北パターンの2つです．連勝数がX以上になるパターンは，考慮不要です．

X=3，A=[2, 3]の例

bは以下のように推移します．

• b: ((0, False), 1)
• b: ((2, False), 1), ((0, False), 2)
• b: ((0, True), 1), ((0, False), 7), ((3, False), 2)

1ステージ終了後の全パターン数は，1 * (A[0] + 1) = 3です．
2ステージ終了後の全パターン数は，3 * (A[1] + 1) = 12です．
bには4連勝，5連勝したパターンが記録されない点に注意が必要です．
最大X=3連勝するパターン数は，3になります．

終わりに

Rank Sの問題は，素のPythonで解くのは大変そうでした．
一方で，適切なパッケージや関数を知っていると楽に解ける問題も多かったです．
最後の問題「十億連勝」は，楽に解く方法が見つからず，一番苦労しました．

関連記事