ABC202A~DをPythonで解いてみた。

エイシングプログラミングコンテスト2021（ABC202)

A.Three Dice
B.180°
C.MadeUp
D.aab aba baa
最後に

問題等引用元：エイシングプログラミングコンテスト2021
https://atcoder.jp/contests/abc202

使っている入力コード

import sys
readline = sys.stdin.readline
readall = sys.stdin.read
ns = lambda: readline().rstrip()
ni = lambda: int(readline().rstrip())
nm = lambda: map(int, readline().split())
nl = lambda: list(map(int, readline().split()))
prl = lambda x: print(*x ,sep='\n')
string = 'abcdefghijklmnopqrstuvwxyz'

A

問題

高橋君が 3つのサイコロを振ったところ、出た目はそれぞれa,b,cでした。これらのサイコロについて、出た目とは反対の面が表す整数を足し合わせた値を求めてください。

解答

a,b,cそれぞれについて7から引いた数が反対の面の数なので、その和を求めれば良いです。

main.py

l = nl()
res = 0
for i in range(3):
    res += 7-l[i]
print(res)

B

問題

0,1,6,8,9からなる文字列Sが与えられます。
Sを180度回転したものを出力してください。

解答

紙に文字列Sを書いて紙を180度回転させたらどのような文字ができるかという問題ですが、変換規則は

・Sを反転する。
・0→0、1→1、6→9、8→8、9→6に変換する
となると問題文で書かれているのでそのままコードを書けばよいです。

main.py

s = list(ns())
s = reversed(s)
dic = {'0':'0','1':'1','6':'9','9':'6','8':'8'}
res = []
for i in s:
    res.append(dic[i])
print(''.join(res))

C

問題

1以上N以下の整数からなる長さNの数列
A={A_i}、B={B_i}、C={C_i}   (i=1,2...N)  があり、1以上N以下の整数i、jの組(i,j)についてA_i=B_Cj となる総数を求めてください。

解答

A、Bが持つ数字ごとにindexをメモする。
数字ごとにA、Bにその数字があるかをチェックし、あれば先程メモしたその数字になるBのindexがCにあるかをチェックする。
(Aのindexの数)×(Cのindexの数)の和を計算すれば完了。

main.py

n = ni()
a = nl()
b = nl()
c = nl()
adic = {i:[] for i in range(1,n+1)}
bdic = {i:[] for i in range(1,n+1)}
cdic = {i:[] for i in range(1,n+1)}
for i in range(n):
    adic[a[i]].append(i+1)
    bdic[b[i]].append(i+1)
    cdic[c[i]].append(i+1)
res = 0
check = []
for i in range(1,n+1):
    l1 = adic[i]
    l2 = bdic[i]
    if len(l1) == 0 or len(l2) == 0:
        continue
    check.append(i)
for i in check:
    for j in bdic[i]:
        l3 = cdic[j]
        if len(l3) == 0:
            continue
        res += len(adic[i]) * len(l3)
print(res)

D

問題

A個のaとB個のbからなる長さA+Bの文字列のうち、辞書順でK番目のものを求めてください。

解答

求めたい文字列をS(a,b,k)、(a,b)が作る文字列の個数をC(a,b)とします。
例えばa=3,b=3,k=5を考えます。
順番に書いてみるとaaabbb,aababb,aabbab,aabbba,abaabbとなりますが、これをどうやってつくればよいでしょうか。6文字は5文字にa,bをつけたもので5文字は4文字の...と考えればよさそうです。そこで、a、ｂどちらをつければよいかを考えます。
先程列挙した文字列の先頭のaを除いたaabbb,ababb,abbab,baabb,...は先ほどと変わらず1,2,3...,C(a-1,b)番目になっています。この中にk番目が入っていれば、aを先頭に加えて5文字のk番目は6文字のk番目となります。反対にaを除いた文字列にk番目が入っていない場合は先頭にbを加えることでC(a-1,b)+1,...,k,...,C(3,3)を調べられます。つまり、C(a-1,b)>=kであれば先頭にaをつける。C(a-1,b)<kであればbを先頭につけるのを繰り返します。ただし、bをつけたときには、k'=k-C(a-1,b)としなければならないことに注意しましょう。なぜなら、bを先頭に付けた文字列からは残ったk-C(a-1,b)番目の文字列を探すからです。以上からS(a,b,k)は再帰関数を使えば実現できます。
また、C(a,b)は動的計画法を用いれば計算できます。動的計画法と言っても難しくはなく、(a+1)×(b+1)のマスを用意((0,0)には1をセット)。(0,0)からスタートして(i,j)には(i-1,j)までの個数+(i,j-1)までの個数を入れるとよいです。以上を実装すると下のようになります。

main.py

def saiki(an,bn,c,k):
    if an == 0:
        return 'b'*bn
    if bn == 0:
        return 'a'*an
    if c[an-1][bn] >= k:
        return 'a' + saiki(an-1,bn,c,k) 
    if c[an-1][bn] < k:
        return 'b' + saiki(an,bn-1,c,k-c[an-1][bn])
    
a,b,k = nl()
c = [[0]*(b+1) for _ in range(a+1)]
c[0][0] = 1
for i in range(a+1):
    for j in range(b+1):
        if i > 0:
            c[i][j] += c[i-1][j]
        if j > 0:
            c[i][j] += c[i][j-1]
res = saiki(a,b,c,k)
print(res)