AtCoder Beginner Contest 328の振り返り

結果はこちらです。

A - Not Too Hard

この問題はX以下の数の和を出力すれば大丈夫です。

N,X=map(int,input().split())
S=list(map(int,input().split()))
print(sum(i for i in S if i<=X))

B - 11/11

この問題はiに対してj(1<=j<=D[i])としてすべて同じ文字になればカウントをインクリメントすればいいです。

N=int(input())
D=[0]+list(map(int,input().split()))
ans=0
for i in range(1,N+1):
  for j in range(1,D[i]+1):
    subs=set()
    for k in range(len(str(i))):subs.add(str(i)[k])
    for k in range(len(str(j))):subs.add(str(j)[k])
    if(len(subs)==1):ans+=1
print(ans)

C - Consecutive

この問題は累積和を用いて解くことができます。あるところまで何回同じ英小文字が連続していたかの回数を累積することでLからRまでの回数がO(1)で出せます。
またわかりやすく1-indexedにしています。

N,Q=map(int,input().split())
S='_'+input().rstrip()
s=[0]
for i in range(1,N+1):
  c=S[i]==S[i-1]
  s.append(s[-1] + c)
for i in range(Q):
  l,r=map(int,input().split())
  print(s[r]-s[l])

D - Take ABC

この問題はstackを用いて解くことができます。空文字Sに対して文字を追加していって最後の末尾がABCになるのであればpopを3回行ってあげます。今回はdequeを用いて解きました。

from collections import deque
S=input().rstrip()
deq=deque()
for i in range(len(S)):
  deq.append(S[i])
  if(len(deq)>=3):
    if(deq[-1]=='C' and deq[-2]=='B' and deq[-3]=='A'):
      for _ in range(3):deq.pop()
print("".join(deq))

E - Modulo MST

この問題は制約から全ての全域木を調べてそのコストが最小のものを出力すればいいことがわかります。またこの問題の場合N頂点あるのでN-1本の辺を全て張り、それが全域木なのであればコストを最小に更新、という方針を立てて解くことができます。

unionfindを用いて全域木かどうか判断する方法

全域木(Spanning Tree)とは、与えられた無向グラフの頂点を全て含み、サイクルを持たないグラフのことです。
unionfindで全域木かどうかの判断をするには、
1.サイクルがあるかどうか、
2.辺の数が頂点数N-1かどうか、
3.グループの数が1かどうか、
で判断することができます。今回の問題の場合は、辺の数はN-1で固定してるので判断しなくて大丈夫です。またN-1本の辺がサイクルがなく張れるのであればグループの数は1となるので判断しなくて大丈夫です。よってサイクルがあるかどうか、だけを判断してあげれば大丈夫です。

class UnionFind:
    def __init__(self, n):
        self.parent = [-1] * n

    def find(self, x):
        if self.parent[x] < 0:
            return x
        else:
            self.parent[x] = self.find(self.parent[x])
            return self.parent[x]

    def merge(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x != y:
            if self.parent[x] > self.parent[y]:
                x, y = y, x
            self.parent[x] += self.parent[y]
            self.parent[y] = x

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def size(self, x):
        return -self.parent[self.find(x)]

from itertools import combinations
N,M,K=map(int,input().split())
g=[]
for _ in range(M):
  u,v,w=map(int,input().split())
  u-=1
  v-=1
  g.append([u,v,w%K])
ans=float('inf')
for comb in combinations(g,N-1):
  uf=UnionFind(N)
  ok=1
  tmp=0
  for u,v,w in comb:
    if(uf.same(u,v)):ok=0
    else:
      uf.merge(u,v)
      tmp+=w
      tmp%=K
  if ok:ans=min(ans,tmp)
print(ans)

F - Good Set Query

この問題はポテンシャル付きunionfindで解くことができます。マージできるかどうかを判断するだけで解くことができます。

class WeightedUnionFind:
    def __init__(self, n):
        self.par = [i for i in range(n+1)]
        self.rank = [0] * (n+1)
        # 根への距離を管理
        self.weight = [0] * (n+1)

    # 検索
    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            y = self.find(self.par[x])
            # 親への重みを追加しながら根まで走査
            self.weight[x] += self.weight[self.par[x]]
            self.par[x] = y
            return y

    def merge(self, x, y, w):
        rx = self.find(x)
        ry = self.find(y)

        # 既に同じ集合に属している場合、重みの矛盾をチェック
        if rx == ry:
            # 重みの矛盾がある場合はFalseを返す
            if self.diff(x, y) != w:
                return False
            return True

        # 木の高さに基づく併合
        if self.rank[rx] < self.rank[ry]:
            self.par[rx] = ry
            self.weight[rx] = w - self.weight[x] + self.weight[y]
        else:
            self.par[ry] = rx
            self.weight[ry] = -w - self.weight[y] + self.weight[x]
            # 木の高さが同じだった場合の処理
            if self.rank[rx] == self.rank[ry]:
                self.rank[rx] += 1
        return True

    # 同じ集合に属するか
    def same(self, x, y):
        return self.find(x) == self.find(y)

    # xからyへのコスト
    def diff(self, x, y):
        return self.weight[x] - self.weight[y]


N,Q=map(int,input().split())
ans=[]
puf=WeightedUnionFind(N)
for i in range(Q):
    a,b,d=map(int,input().split())
    a-=1
    b-=1
    if(puf.merge(a,b,d)):ans.append(i+1)
print(*ans)

まとめ

今回はD問題で典型のはずなのに時間がかかりすぎて焦りました、、。ただE,F問題を解くことができたので、しっかりと問題を典型に落とし込んで解けるようにしていきたいと思っています。