from collections import defaultdict

N, K = map(int, input().split())
A = list(map(int, input().split()))
if K == 1:
    print("Yes")
    exit()
dict_set = defaultdict(lambda: defaultdict(int))
for i, a in enumerate(A):
    dict_set[i % K][a] += 1
A.sort()
for i, a in enumerate(A):
    if dict_set[i % K][a] > 0:
        dict_set[i % K][a] -= 1
        # NOTE: 試験中はこれで出したが、実際には不要な部分をコメントアウト
        # if dict_set[i % K][a] <= 0:
        #   del dict_set[i % K][a]
    else:
        print("No")
        exit()
print("Yes")

Submission #32223209 - AtCoder Beginner Contest 254

ポイントは下記だと考えます。個人的には、C問題だけど、結構難しかったなと思いました。

K=1の時は、絶対sortできるのでTrue
K>=2 の時は、交換できない項が発生するので、交換できる項のグループ(= index番号をKで割った余り)をkeyとしたdefaultdictで保持し、重複もするため、どの数字が何個あるかという情報を保持する(= defaultdict(lambda: defaultdict(int)))
Aをsortした数列を用意し、1つづつ、該当のグループ内に数値があるかを確認する、なければFalse、あればそのグループ内の数字のcntを減らし、次に進む
全て完了すればTrueとなる

自分は、sortしたAと比較しましたが、グループ分けしたそれぞれの配列をsortしてみるという方法でもできます。下記の解説が詳しいです。

www.youtube.com

E問題

atcoder.jp

E問題は、試験中取り組まなかったのですが、解説見ずに挑戦し解くことができました。ポイントは、下記だと考えます。

kが3以下と少ない
出次数が3以下が保証されいている

そのため、下記の図のように探索範囲が狭い(3通り * 2通り * 2通り= 12通り)ことがわかります。

そのため、隣接リスト等でグラフ構造を変数で保持できて、BFSが使えれば、TLEせずに解くことができます。水色diffですが、一見すると難しそうですが、上記に気付けさえすれば、正直茶色レベルでもおかしくないのかな？と思うような問題でした。

from collections import deque

N, M = map(int, input().split())
edge = [[] for _ in range(N)]
for _ in range(M):
    a, b = map(int, input().split())
    edge[a-1].append(b-1)
    edge[b-1].append(a-1)
ans = []
Q = int(input())
for _ in range(Q):
    x, k = map(int, input().split())
    cnt, finds, queue = x, set([x]), deque([(x-1, 0)])
    while queue:
        cur, dis = queue.popleft()
        for nxt in edge[cur]:
            if dis + 1 > k or nxt + 1 in finds:
                continue
            cnt += nxt + 1
            finds.add(nxt+1)
            queue.append((nxt, dis + 1))
    ans.append(cnt)
print(*ans, sep="\n")

Submission #32247915 - AtCoder Beginner Contest 254

D問題

atcoder.jp

個人的には、解けそうと思って勘違いあって、時間潰してしまいましたが、結局解説見ないと解けなかったです。シンプルだけど整数苦手でこの手の問題落としがちだなぁと思っています。

まず、i, j の全探索は上記の通りでTLEするので、どのようにして効率よく考えられるか？を考えます。そこで、平方数は 5 *5 みたいに(1, 25), (25, 1), (5, 5) の3通りで、それぞれN以下のもの探せば良さそうと思ったのですが。この時、36みたいなパターンもあることを抜けてサンプル通らず、実際全探索で書いてみて、2² * 3³ みたいなの、どう数え上げよう。。。と思って時間切れになりました。

atcoder.jp

www.youtube.com

公式解説や、かつっぱさんの動画でも解説されていますが、平方数とはなんぞや？何と何書けたら平方数になるの？というところがポイントな気がしています。

解説だとわかりにくいのですが、要は平方数は素因数分解した際に、すべての次数が偶数になるものである。そして、i*jが平方数になる場合は、i, jをそれぞれ因数分解した際にできる最大の平方数(=解説では、約数の内最大の平方数といっている)と、それ以外と考えた時に、それ以外のもの同志を掛け合わせて平方数にする必要がある。ということは、それぞれは同じ数でないといけない！ということがわかると思います。

from collections import defaultdict

N = int(input())
cnt = defaultdict(int)
for i in range(1, N+1):
    j = 2
    while i > 1 and i >= j*j:
        while i % (j*j) == 0:
            i //= j*j
        j += 1
    cnt[i] += 1
ans = 0
for k in cnt:
    ans += cnt[k]*cnt[k]
print(ans)