第三回 アルゴリズム実技検定

 A ~ O の全15問中、Gまで7問完答で初級でした。普段のAtCoderと違って実務よりの問題が多いとは聞いてましたが、確かに現実にありそうなシーンを意識して作問してる雰囲気は感じました。今回のAtCoderアルゴリズム実技検定はコロナウイルス影響で無料受験可能だったので、これを機に初めて受けてみたのですが、やはりもっと基本的なとこから鍛えていかなければならないなと改めて実感。

G問題までの解法

 自分の解法。もっと効率的な方法はあると思う。

種類 概要
A str.lower() メソッドで
B 与えられたクエリを一つ一つ実直に
C そのままfor文。ただし R=1の時はすぐAを出力。じゃないと R=1の時に無限ループになる。
D めちゃくちゃ大変。if文の列挙。解答例は割と上手くやっていた(それでもややこしい)
E 無効グラフの問題。制約的にそのまま実直にクエリを実行
F 最初問題文の意味が分からずかなり時間取られる。上手い方法が思いつかずqueを使って実直に実装
G 幅優先探索



H - ハードル走

問題のリンク

概要

  • 3つの行動から一つ選んで実行できる。(詳細は問題リンクから)
  •  0からスタートして、 Lまでいく時の最短時間を求める。
  • ただし間に N個のハードルがある。

制約

  •  2 ≤ L ≤ 10^{5}
  •  1 ≤ N ≤ L
  •  0 \lt x_1 \lt x_2 \lt \cdots \lt x_N \lt L
  •  2 ≤ T_1, T_2, T_3 ≤ 1000
  •  T_1, T_2, T_3は偶数

考察

 問題の感じからDP。もしくは貪欲法あたりだろうなと考えてました。どうやらDPで行くらしいです。考えた方は以下の通り。

  • dp[0, inf, inf, inf, ... ]で初期化。
  • min(dp[i+1], dp[i]+行動1でかかる時間)で更新。もし座標 iにハードルがあった場合T3を後者に追加
  • min(dp[i+2], dp[i]+行動2でかかる時間)で更新。もし座標 iにハードルがあった場合T3を後者に追加
  • min(dp[i+4], dp[i]+行動3でかかる時間)で更新。もし座標 iにハードルがあった場合T3を後者に追加

 基本は上記の様にして動的計画法( Order(N))してく感じ。ただゴール近くに差し迫った場合、ジャンプ中に飛び越している可能性もあるので、最後に微調整が必要。

コード

N, L = map(int, input().split())
X = [False] * (L+10)
hard = list(map(int, input().split()))

for x in hard:
    X[x] = True

T1, T2, T3 = map(int, input().split())

dp = [float('inf')]*(L+10)
dp[0] = 0

for i in range(L+1):
    cost = T3 if X[i] else 0
    dp[i+1] = min(dp[i+1], dp[i] + T1 + cost)
    dp[i+2] = min(dp[i+2], dp[i] + T1 + T2 + cost)
    dp[i+4] = min(dp[i+4], dp[i] + T1 + T2*3 + cost)
    if i+1 == L:
        dp[i+1] = min(dp[i+1], dp[i] + T1//2 + T2//2 + cost)
    elif i+2 == L:
        dp[i+2] = min(dp[i+2], dp[i] + T1//2 + 3*T2//2 + cost)
    elif i+3 == L:
        dp[i+3] = min(dp[i+3], dp[i] + T1//2 + 5*T2//2 + cost)
 
print(dp[L])



I - 行列操作

問題のリンク

考察

 普通に実装すると時間計算量以前に空間計算量的にもヤバくてメモリエラー起こしてました。解法みた所、実はそれに近い形で頭の中では解き方をイメージしてたものの、実装ミスで解けなかった感じです。

コード

N = int(input())
rows = [i for i in range(N)]
columns = [i for i in range(N)]
transposed = False

Q = int(input())
for _ in range(Q):
    q = list(map(int,input().split()))

    if q[0] == 3:
        transposed = not transposed
    else:
        A, B = q[1]-1, q[2]-1
        if q[0] == 1:
            if transposed:
                columns[A], columns[B] = columns[B], columns[A]
            else:
                rows[A], rows[B] = rows[B], rows[A]
        elif q[0] == 2:
            if transposed:
                rows[A], rows[B] = rows[B], rows[A]
            else:
                columns[A], columns[B] = columns[B], columns[A]
        elif q[0] == 4:
            if transposed:
                A, B = B, A
            r,c = rows[A], columns[B]

            print(N*r+c)


PyTorch Scheduler Class の Chaining

今回はPyTorch 1.4.0から追加された新機能であるSchedulerのChainingについてです。この方のQiita記事で知ったのですが、PyTorchのschedulerに新機能が追加された様です。

qiita.com

内容はシンプルで、今までは学習ループの中に1種類しか組み込めなかったschedulerを複数重ねて使えると言うことです。具体例をみていきましょう。

import torch
from torch.optim.lr_scheduler import ExponentialLR, CosineAnnealingLR
from torch.optim import SGD
import matplotlib.pyplot as plt
import seaborn as sns
sns.set()

lr_cos, lr_exp, lr = [], [], []

# 初期化
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 1e-2)
optimizer.zero_grad()
scheduler_cos = CosineAnnealingLR(optimizer, T_max=20, eta_min=1e-6)
for epoch in range(200):
    optimizer.step()
    optimizer.zero_grad()
    scheduler_cos.step()
    lr_cos.append(scheduler_cos.get_last_lr()[0])

# 初期化
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 1e-2)
optimizer.zero_grad()
scheduler_exp = ExponentialLR(optimizer, gamma=0.97)
for epoch in range(200):
    optimizer.step()
    optimizer.zero_grad()
    scheduler_exp.step()
    lr_exp.append(scheduler_exp.get_last_lr()[0])

# 初期化
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 1e-2)
optimizer.zero_grad()
scheduler_cos = CosineAnnealingLR(optimizer, T_max=20, eta_min=1e-6)
scheduler_exp = ExponentialLR(optimizer, gamma=0.97)
for epoch in range(200):
    optimizer.step()
    optimizer.zero_grad()
    scheduler_exp.step()
    scheduler_cos.step()
    lr.append(optimizer.param_groups[0]['lr'])

plt.plot(lr_cos, label='CosineAnnealing', linestyle='dashed')
plt.plot(lr_exp, label='Exponential', linestyle='dashed')
plt.plot(lr, label='Fusion', color='red', linewidth=3)
plt.legend(bbox_to_anchor=(0, -0.1), loc='upper left', borderaxespad=0)
plt.show()

f:id:spider-man-dance:20200601140933p:plain

この様な形でCosineAnnealingExponentialを使って新たなタイプのschedulerを作れていることが分かります。この機能と以下の方のコードを使えば、今まで自分でクラスを作っていたWarmup系のschedulerも簡単に作れそうです。

github.com

!pip install git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git
from warmup_scheduler import GradualWarmupScheduler

lr = []

# 初期化
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 1e-2)
optimizer.zero_grad()
scheduler_cos = CosineAnnealingLR(optimizer, T_max=400, eta_min=1e-6)
scheduler_wup = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=20, after_scheduler=scheduler_cos)
for epoch in range(200):
    optimizer.step()
    optimizer.zero_grad()
    scheduler_cos.step()
    scheduler_wup.step()
    lr.append(optimizer.param_groups[0]['lr'])

plt.plot(lr, label='Warmup Cosine Annealing', color='red', linewidth=3)
plt.legend(bbox_to_anchor=(0, -0.1), loc='upper left', borderaxespad=0)
plt.show()

f:id:spider-man-dance:20200601141656p:plain


AtCoder ABC 169

ABCD4完の緑パフォでした。今までのABC挑戦の中で、最速でDまで解けたけど(18分)、Cでつまづいてしまいました。考え方については途中気付いたけどなぜか通らずゴニョゴニョしてたら通った感じです。こういったプログラミング言語あるある(桁溢れ問題)は意識していきたい(特に小数扱う際)

A - Multiplication 1

問題のリンク

概要

  •  A×Bを求める。

制約

  •  1 ≤ A, B ≤ 100

考察

自分が見てきた中でAtCoder史上、最も簡単な問題でした。

コード

A, B = map(int, input().split())
print(int(A*B))



B - Multiplication 2

問題のリンク

概要

  •  N個の整数の積を求める。結果が 10^{18}を超える場合は、代わりに-1を出力

制約

  • [tex: 2 ≤ N ≤ 105]
  •  0 ≤ A_i ≤ 10^{18}
  • 入力は全て整数

考察

愚直に実装。0があるかどうかだけ最初に判定してないと、実際の答えは0なのに誤って-1を出力してしまうので注意

コード

N = int(input())
A = list(map(int, input().split()))

if min(A)==0:
    print(0)
    exit()

ans = 1
for a in A:
    ans *= a
    if ans>10**18:
        print(-1)
        exit()

print(ans)



C - Multiplication 3

問題のリンク

概要

  •  A×Bを求める。 Bは小数第2位まで与えられる。

制約

  • Aは整数
  • Bは小数第2位まで与えられる
  •  0 ≤ A ≤ 10^{15}
  •  0 ≤ B ≤ 10

考察

 A問題の小数問題ということで、桁溢れエラーを誘ってる感がぷんぷん漂ってはきていたのですが、結局実装でミスりまくりました。(計7回)  とりあえずなるべく小数問題を避けたいので、文字列やら整数やらで代用できる様にゴニョゴニョしてたらいつの間にか解けたって感じですが、次からはこういうの気をつけたい。

コード

from math import floor

A, B = input().split()
A = int(A)

B1 = int(B[:-3])
B2 = int(str(B)[-2:])

ans1 = A * B1
ans2 = A * B2 // 100

ans = ans1 + ans2
print(floor(ans))



D - Div Game

問題のリンク

概要

  • はじめに整数 Nが与えられる。 Nに対して、以下の操作を何度できるか。
- ある素数[tex: p]と正の整数[tex: e]を用いて、[tex: z = p^e]と表せる。
- [tex: N]が[tex: z]で割り切れる。
- 以前の操作で選んだどの整数とも異なる。
- [tex: N]を[tex: N/z]に置き換える。

制約

  • 入力は全て整数
  •  1 ≤ N ≤ 10^{12}

考察

  • まず、 N素因数分解
  • 次にそれぞれの素因数から問題文の条件を満たす zがどれだけ作れるかをループでカウント

コード

N = int(input())

def factorization(n):
    arr, temp = [], n
    for i in range(2, int(-(-n**0.5//1))+1):
        if temp%i==0:
            cnt=0
            while temp%i==0:
                cnt+=1
                temp //= i
            arr.append([i, cnt])
    if temp!=1:
        arr.append([temp, 1])
    if arr==[]:
        arr.append([n, 1])
    return arr

fac = factorization(N)
ans = 0

for k, v in fac:
    cnt = 1
    if k != 1:
        while True:
            v -= cnt
            if v < 0:
                break
            ans += 1
            cnt += 1

print(ans)



AtCoder ABC 168

ABC3完の茶色パフォでした。今回のレベルでDまで行けなかったのは不甲斐ないので、改めて精進したい。

A - ∴ (Therefore)

問題のリンク

概要

  •  Nの一の位で場合分け

制約

  •  N 999以下の正の整数

考察

特になし

コード

N = input()

if int(N[-1]) in [2,4,5,7,9]:
    print('hon')
elif int(N[-1]) in [0,1,6,8]:
    print('pon')
else:
    print('bon')



B - ... (Triple Dots)

問題のリンク

概要

  • 文字列 Sの長さで場合分け

制約

  •  K 1以上、 100以下の整数
  •  Sは英小文字からなる文字列
  •  Sの長さは 1以上、 100以下

考察

特になし

コード

K = int(input())
S = input()

if len(S)<=K:
    print(S)
else:
    ans = S[:K] + '...'
    print(ans)



C - : (Colon)

問題のリンク

概要

  • ある時計の長針の長さ、短針の長さ、及び現時刻が与えられた時の長針の先端と短針の先端の距離

制約

  • 入力は全て整数
  •  1 ≤ A, B ≤ 1000
  •  0 ≤ H ≤ 11
  •  0 ≤ M ≤ 59

考察

 三平方の定理をググれば理論上はすぐに解ける問題。だが当初、短針の角度を(H/12) * 360)としてしまっており、計算が合わず時間をロスしてしまった。正しくはminuteも考慮した((H/12) * 360) + (30*(M/60))。これに気づくまでに時間がかかってしまった。今後は気を付けたい。

コード

import math

A, B, H, M = map(int, input().split())

ka_a = ((H/12) * 360) + (30*(M/60))
ka_b = M/60 * 360

kakudo = max(ka_a, ka_b) - min(ka_a, ka_b)
kakudo = min(kakudo, (360-kakudo))
cos = math.cos(math.radians(kakudo))

if kakudo==180:
    print(A+B)
    exit()

ans = A**2 + B**2 - 2*A*B*cos

print(math.sqrt(ans))



D - .. (Double Dots)

問題のリンク

概要

  • いわゆる最短経路問題
  • 最短経路を通った際、次に進む先のノードの番号を求める

制約

  • 入力は全て整数
  •  2 ≤ N ≤ 10^{5}
  •  1 ≤ M ≤ 2×10^{5}
  •  1 ≤ A_i, B_i ≤ N(1 ≤ i ≤ M)
  •  A_i ≠ B_i(1 ≤ i ≤ M)
  • どの 2部屋間も、道路をいくつか通って行き来できる

考察

  • 根( 1番目の部屋)をスタートに幅優先探索をしていく
  • 制約上、どの二部屋間も行き来できる(つまり必ず1のゴールにたどり着けるという保証がされている)ので、Noの答えはあり得ない。つまり、探索すべきグラフは必ず一つしかあり得ない。

コード

from collections import deque

def bfs_tree(graph, v=1):
    """
    graph: グラフ
    v: スタート
    """
    # 矢印がどこを向いているか
    # -1だと未探索
    arrow = [-1]*len(graph) 
    q = deque()
    q.append(v)
    arrow[v] = 0
    while q:
        now = q.popleft()
        for next_v in graph[now]:
            if arrow[next_v] != -1:
                continue
            else:
                q.append(next_v)
                arrow[next_v] = now
    return arrow

N, M = map(int, input().split())
graph = [[] for _ in range(N)]
pre = []

for _ in range(M):
    A, B = map(int, input().split())
    graph[A-1].append(B-1)
    graph[B-1].append(A-1)

arrow = bfs_tree(graph, 0)

print('Yes')
for i in range(1, N):
    print(arrow[i]+1)



AtCoder ABC 167

ABCD4完の緑パフォでした。

A - Registration

問題のリンク

概要

  • 以下の二つの文字列が等しいか否か
  •  S, 及び Tの末尾を削ったもの

制約

  •  1 ≤ |S| ≤ 10
  •  |T| = |S| + 1

考察

特になし

コード

S = input()
T = input()

if S == T[:-1]:
    print('Yes')
else:
    print('No')



B - Easy Linear Programming

問題のリンク

概要

  •  1が書かれたカードが  A枚、 0が書かれたカードが B枚、 −1が書かれたカードが C枚あります。これらのカードから、ちょうど K枚を選んで取るとき、取ったカードに書かれた数の和として、 ありうる値の最大値はいくつか。

制約

  • 入力は全て整数である。
  •  0 ≤ A, B, C
  •  1 ≤ K ≤ A + B + C ≤ 2×10 ^{9}

考察

特になし

コード

A, B, C, K = map(int, input().split())

if K<=A:
    print(K)
elif A<K<=A+B:
    print(A)
elif A+B<K<=A+B+C:
    print(A-(K-A-B))



C - Skill Up

問題のリンク

概要

  • 高橋くんは学びたいアルゴリズム M個あり、スタート時点で M個全て理解度は 0である。
  • 参考書が N冊ある。 i番目の参考書を読むとそれぞれのアルゴリズムの理解度が A_{ij}上がる。また i番目の参考書は C_i円である。
  • 全ての理解度が X以上になるための最小コストを求めよ。

制約

  • 入力は全て整数
  •  1 ≤ N, M ≤ 12
  •  1 ≤ X ≤ 10 ^{5}
  •  1 ≤ C_i ≤ 10 ^{5}
  •  1 ≤ A_{ij} ≤ 10 ^{5}

考察

 とりあえず全探索が間に合いそうなことは分かった。(解説でも全探索しか方法はないと言っていた。)そこで買う or 買わないのbit全探索を実装。ただ、少しコードが長くなって正解するまでに実装ミスを連発してしまった。尚、DFSで実装する方が効率的。

コード

import numpy as np

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

ans = 10**9

for i in range(1<<N):
    cost = 0
    tisiki = np.array([0]*M)
    cond = [0]*N
    for j in range(N):
        if 1&(i>>j):  # 右シフトして論理積(bit全探索)
            cond[j] = 1
    
    for k in range(N):
        if cond[k]==1:
            b = book[k]
            cost += b[0]
            tisiki += np.array(b[1:])
    
    if np.all(tisiki>=X):
        ans = min(ans, cost)

if ans < 10**9:
    print(ans)
else:
    print(-1)



D - Teleporter

問題のリンク

概要

  •  1 ~  Nまでランダムに番号のついた N個の街がある。
  • それぞれの街にはテレポーテーターが設置されており、それぞれの転送先は A_iである。
  •  K回移動した際に泊まっている街を求めよ。

制約

  •  2 ≤ N ≤ 2×10^{5}
  •  1 ≤ A_i ≤ N
  •  1 ≤ K ≤ 10 ^{18}

考察

  •  Kが物凄く大きいためそのまま Kfor文を回すのは不可能。
  • 一度ループが検出できれば、残っている移動回数から最終的なストップ地点を求めることが可能。
  • 移動していく処理は再帰関数で実装。
  • 他の人のコードを見てるとどうやら再帰関数を使わずwhile文を使って書いてる人もいた。

コード

import sys
sys.setrecursionlimit(10**9)

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

log = []
check = [False]*N

def func(x, iter):
    x = A[x-1]

    if iter+1 == K:
        return x
    if check[x-1]:
        ix = log.index(x)
        loop = log[ix:]
        tmp = (K-iter-1)%len(loop)
        return loop[tmp]
    else:
        log.append(x)
        check[x-1] = True
        iter += 1
        return func(x, iter)

print(func(1, 0))



AtCoder ABC 165

ABCD4完で水色パフォでした。以下振り返り。

A - We Love Golf

問題のリンク

概要

  • 要は A以上、 B以下の数に Kの倍数が含まれるか否か。

制約

  •  1 ≤ A ≤ B ≤ 1000
  •  1 ≤ K ≤ 1000

考察

面倒なので愚直に全探索

コード

K = int(input())
A, B = map(int, input().split())

for i in range(A, B+1):
    if i%K==0:
        print('OK')
        exit()

print('NG')



B - 1%

問題のリンク

概要

  • 初年度に100円を預ける。年率1%の利子がつく時、 X以上になるのは何年後か。

制約

  •  101 ≤ X ≤ 10 ^{18}

考察

最初、1年毎に端数切り捨てという条件を忘れて実装してしまってたけど、提出前に確認して気づけた。(本番慣れしてきた証拠と思いたい。)

コード

X = int(input())

cnt = 0
m = 100
while True:
    m *= 1.01
    m = int(m)
    cnt += 1
    if m>=X:
        break

print(cnt)



C - Many Requirements

問題のリンク

概要

文字にするとややこしいので省略。

制約

ながいので省略。ただ制約的に全探索でも間に合う感じ。

考察

  • 問題文を始め読んだ際、日本語の意味が理解できず5分以上は考え込んでしまった。
  • 理解してしまえば、そのまま全探索で行ける問題。
  • ただ初回の提出の際にcombinations_with_replacement(range(1, M+1), N)combinations_with_replacement(range( M+1), N)と書いてしまい、0を含めた数列で探索してしまった。(1WA)
  • C問題はBまでと比べると実装が少し重くなり始めるけど、こういう系はbit全探索系の問題を何度か解く事で慣れてきた気がする。

コード

from itertools import combinations_with_replacement

N, M, Q = map(int, input().split())
q = []
for _ in range(Q):
    q.append(list(map(int, input().split())))

ans = 0
for v in combinations_with_replacement(range(1, M+1), N):
    score = 0
    for i in range(Q):
        joken = q[i]
        if v[joken[1]-1] - v[joken[0]-1] == joken[2]:
            score += joken[3]
    ans = max(ans, score)
    #print(v, ': ', score)

print(ans)



D - Floor Function

問題のリンク

概要

  •  N以下の非負整数 xに対する floor(Ax/B) -A × floor(x/B) の最大値

制約

  •  1 ≤ A ≤ 10 ^{6}
  •  1 ≤ B ≤ 10 ^{12}
  •  1 ≤ N ≤ 10 ^{12}
  • 入力は全て整数

考察

  •  floor(x/B) が0となる範囲内で、最大の xを求める。要は x=B-1
  • ただ xが指定の範囲を取れない場合もある。よってその時の条件を忘れないことが肝。( x=N)

コード

from math import floor

A, B, N = map(int, input().split())

if N>=B-1:
    x = B-1
    ans = floor((A*x)/B) - A*floor(x/B)
else:
    x = N
    ans = floor((A*x)/B) - A*floor(x/B)

print(ans)



E - Rotation Matching

問題のリンク

今回、できなかった問題。テストケース8割は通ったけど、残りがダメでした。

概要

  • 参加者 N人でジャンケン大会を行う。参加者はそれぞれ 1 ~  Nの番号が割り振られている。
  • 1対1でジャンケンを行う闘技場が M個用意されている。
  • 闘技場には N以下の整数が2個割り当てられている。
  • 参加者は自分の数字が書かれた闘技場でジャンケンを行う(第1ラウンド)
  • その後、参加者に割り振られた整数が1増加する。 N番目の人は1になる。
  • 参加者は自分の数字が書かれた闘技場でジャンケンを行う(第2ラウンド)
  • 以降、 Nラウンドまで繰り返した時に二回以上同じ参加者と対戦するような参加者が存在しないようにしたい。この時の闘技場の整数はどの様に割り当てるか。

制約

  •  1 ≤M
  •  M × 2 + 1 ≤ N ≤ 100000

考察

 公式の解説動画が神すぎるくらい分かりやすかったです。最後におまけで出てくるリーグ戦のアルゴリズムも面白かった!

www.youtube.com



AtCoder ABC 162

ABCD4完の緑パフォでした。以下振り返りと反省。

(2020年4月13日 D問題別解追記)

A - Lucky 7

問題のリンク

概要

  • 3桁の整数 Nが与えられる。
  •  Nのいずれかの桁に7が含まれるか否か。

制約

  •  100 ≤ N ≤ 999
  • 入力は全て整数

考察

特になし

コード

N = input()
if N[0]=='7' or N[1]=='7' or N[2]=='7':
    print('Yes')
else:
    print('No')



B - FizzBuzz Sum

問題のリンク

概要

  •  1から Nまでの整数で、 3でも 5でも割り切れない数を足し合わせる。

制約

  •  1 ≤ N ≤ 10 ^{6}

考察

めちゃくちゃFizzBuzzでした。プログラマーなら誰しもが書いてきたもの。

コード

N = int(input())
ans = 0
for i in range(1, N+1):
    if i%3 != 0 and i%5 != 0:
        ans += i

print(ans)



C - Sum of gcd of Tuples (Easy)

問題のリンク

概要

-  \displaystyle{\sum_{a=1}^K} \displaystyle{\sum_{b=1}^K} \displaystyle{\sum_{c=1}^K}  gcd(a, b, c)を求めよ

制約

  •  1 ≤ K ≤ 200

考察

  • 制約を見るからに a, b, cを全てfor文で探索すると[tex: o(N3)]で間に合わなそう。
  • なので少し工夫をして探索数を減らすことに具体的にはpermutation(順列)ではなく、combinations_with_replacement(重複あり組み合わせ)gcdを求めて行くことに。それで書いたのが以下のような感じ。

コード

from itertools import combinations_with_replacement

from math import gcd   # gcd(6, 4) -> 2
from functools import reduce

def gcd_list(numbers):
    return reduce(gcd, numbers)


N = int(input())
ans = 0
for i, j, k in combinations_with_replacement(range(1, N+1), 3):
    if i==j==k:
        ans += i
    elif i==j:
        ans += gcd(i, k)*3
    elif j==k:
        ans += gcd(j, i)*3
    elif i==k:
        ans += gcd(i, j)*3
    else:
        ans += gcd_list([i, j, k])*6
print(ans)

学び

ちなみに最初は上のコードのmathfractionsにして出したところ普通にTLEでした。。。見た感じギリギリっぽかったのと他に良い案が思い浮かばなかったため、mathに変えたところ普通に通りました。。。時間はロスしたものの、手持ちのスニペットが少し整備されたのでそこは良かった。



D - RGB Triplets

問題のリンク

概要

  • R, G, Bからなる長さ Nの文字列 Sで、以下の条件を満たす (i, j, k)の組みを求めよ。
  •  1 ≤ i \lt j \lt k ≤ N
  •  S_i ≠ S_jかつ S_j≠ S_kかつ S_k ≠ S_i
  •  j-i ≠k-j

制約

  •  1 ≤ N ≤ 4000

考察

  • これも全探索だと O(N^3)でキツそう
  • ijだけの探索なら、 O(N^2)で行けるはず
  • 予め、jの右側にRGBがそれぞれどれだけ存在しているかが分かれば、あとは指定の文字を一々数えなくても自ずと答えはわかる。(累積和を使うことで3つ目のfor文を無くせる。)
  • 例えばi = 'R'j='G'、であれば、jの左側にどれだけBがいくつ存在するかを見れば良い。
  •  j-i ≠k-jの制約だけ注意。もしこれに該当した場合、上記の値から -1

コード

N = int(input())
S = list(input())

num_r = S.count('R')
num_g = S.count('G')
num_b = N - num_r - num_g
C = [[num_r, num_g, num_b] for _ in range(N)]
r_cnt, g_cnt, b_cnt = 0, 0, 0

for i in range(N):
    if S[i] == 'R':
        r_cnt += 1
    elif S[i] == 'G':
        g_cnt += 1
    else:
        b_cnt += 1

    C[i][0] -= r_cnt
    C[i][1] -= g_cnt
    C[i][2] -= b_cnt

ans = 0
for i in range(N-2):
    for j in range(i+1, N-1):
        if S[i] != S[j]:
            if 'R' not in [S[i], S[j]]:
                ans += C[j][0]
                if j+(j-i) < N and S[j+(j-i)]=='R':
                    ans -= 1
            if 'G' not in [S[i], S[j]]:
                ans += C[j][1]
                if j+(j-i) < N and S[j+(j-i)]=='G':
                    ans -= 1
            if 'B' not in [S[i], S[j]]:
                ans += C[j][2]
                if j+(j-i) < N and S[j+(j-i)]=='B':
                    ans -= 1

print(ans)

PyPyで通りました(457ms)。あとで順位表を見てみたところ、このD問題を早く解けるかどうかでだいぶパフォーマンスが違っていたようです。少し実装が複雑になってきたときに素早くバグなく書けるかどうか。。。力の無さを実感します。。。時間出てきたときに後半の問題も追記していきたいと思います。

(2020年4月13日 別解追記) PDF解答をみてなるほどと思ったので、実際にコーディング実装してみました。前のコードに比べるとかなり簡潔に書けてます。ポイントは以下の2点。

  •  SからR G Bの3つのを選ぶ方法はS.count('R') * S.count('G') S.count('B')通り。
  • 制約 j-i ≠k-jを考慮し、この条件に当てはる分を上記の値から引く。ijをカウンターループさせるので O(N^2)
N = int(input())
S = input()

n_R = S.count('R')
n_G = S.count('G')
n_B = N - n_R - n_G
ans = n_R * n_G * n_B

for i in range(N-2):
    for j in range(i+1, N-1):
        k = j + (j-i)
        if k < N:
            if S[i] != S[j] and S[j] != S[k] and S[k] != S[i]:
                ans -= 1

print(ans)

解答見たあとだと、簡単にコーディングもできました。なるほどです。