クルトンのプログラミング教室

Pythonの使い方やPythonを使った競技プログラミングの解法などを解説しています。

Pythonで理解する蟻本「2-6 素数判定」(p.110)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-6 素数判定」(p.110)
のコードをPythonで書き直したものとなっています。

入力


n

入力例



295927


素数判定（O(√n)）

# 素数判定O(√n)
def is_prime(n):
    i = 2
    while i * i <= n:
        if n % i == 0:
            return False
        i += 1
    return n != 1    # 1の場合は例外

実行結果

False

約数の列挙（O(√n)）

# 約数の列挙O(√n)
def divisor(n):
    res = []
    i = 1
    while i * i <= n:
        if n % i == 0:
            res.append(i)
            if i != n // i:
                res.append(n // i)
        i += 1
    return res

実行結果

[1, 295927, 541, 547]

素因数分解（O(√n)）

# 素因数分解O(√n)
def prime_factor(n):
    res = {}
    i = 2
    while i * i <= n:
        while n % i == 0:
            if not i in res:
                res[i] = 0
            res[i] += 1
            n //= i
        i += 1
    if n != 1:
        res[n] = 1
    return res

実行結果

{541: 1, 547: 1}

Pythonで理解する蟻本「2-6 双六」(p.108)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-6 双六」(p.108)
のコードをPythonで書き直したものとなっています。

入力


a\,b

入力例



4 11


拡張ユークリッドの互除法

def extgcd(a, b):
    d = a
    if b != 0:
        d, y, x = extgcd(b, a % b)
        y -= (a // b) * x
    else:
        x = 1
        y = 0
    return d, x, y

解答

# 入力
a, b = map(int,input().split())

# 拡張ユークリッドの互除法
def extgcd(a, b):
    d = a
    if b != 0:
        d, y, x = extgcd(b, a % b)
        y -= (a // b) * x
    else:
        x = 1
        y = 0
    return d, x, y

d, x, y = extgcd(a, b)

if d == 1:
    L = [0] * 4
    if x > 0:
        L[0] += x
    else:
        L[2] += -x
    if y > 0:
        L[1] += y
    else:
        L[3] += -y
    print(*L)
else:
    print(-1)

Pythonで理解する蟻本「2-6 線分上の格子点の個数」(p.107)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-6 線分上の格子点の個数」(p.107)
のコードをPythonで書き直したものとなっています。

入力


x_1\,y_1\\x_2\,y_2

入力例



1 11
5 3


最大公約数を求める関数

def gcd(a, b):
    if b == 0:
        return a
    return gcd(b, a % b)

解答

# 入力
x1, y1 = map(int,input().split())
x2, y2 = map(int,input().split())

# 最大公約数を求める関数
def gcd(a, b):
    if b == 0:
        return a
    return gcd(b, a % b)

ans = gcd(abs(x1 - x2), abs(y1 - y2)) - 1
print(ans)

Pythonで理解する蟻本「2-5 Layout（POJ No.3169）」(p.104)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 Layout（POJ No.3169）」(p.104)
のコードをPythonで書き直したものとなっています。

入力


N\,ML\,MD\\AL_1\,BL_1\,DL_1\\:\\AL_{ML}\,BL_{ML}\,DL_{ML}\\AD_1\,BD_1\,DD_1\\:\\AD_{MD}\,BD_{MD}\,DD_{MD}

入力例



4 2 1
1 3 10
2 4 20
2 3 3


解答

import sys

MAX_N = 1000
MAX_ML = 10000
MAX_MD = 10000
INF = float('inf')

# 入力
N, ML, MD = map(int,input().split())
AL = [0] * MAX_ML
BL = [0] * MAX_ML
DL = [0] * MAX_ML
for i in range(ML):
    AL[i], BL[i], DL[i] = map(int,input().split())
AD = [0] * MAX_MD
BD = [0] * MAX_MD
DD = [0] * MAX_MD
for i in range(MD):
    AD[i], BD[i], DD[i] = map(int,input().split())

d = [0] * MAX_N    # 最短距離
updated = False   # 更新されたか

def update(x, y, i):
    global updated
    if x > y:
        d[i] = y
        updated = True

# ベルマンフォード法によりdを計算する
def bellmanford():
    global updated
    for k in range(N + 1):
        updated = False
        # i+1からiへコスト0
        for i in range(N):
            if d[i + 1] < INF:
                update(d[i], d[i + 1], i)
        # ALからBLへコストDL
        for i in range(ML):
            if d[AL[i] - 1] < INF:
                update(d[BL[i] - 1], d[AL[i] - 1] + DL[i], BL[i] - 1)
        # BDからADへコスト-DD
        for i in range(MD):
            if d[BD[i] - 1] < INF:
                update(d[AD[i] - 1], d[BD[i] - 1] - DD[i], AD[i] - 1)

# 負閉路チェック
bellmanford()
if updated:
    print(-1)
    sys.exit()    # ここでプログラムを終了

d = [INF] * MAX_N
d[0] = 0
bellmanford()
res = d[N - 1]
if res == INF:
    res = -2
print(res)

Pythonで理解する蟻本「2-5 Conscription（POJ No.3723）」(p.103)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 Conscription（POJ No.3723）」(p.103)
のコードをPythonで書き直したものとなっています。

入力

N\,M\,R\\x_1\,y_1\,d_1\\:\\x_R\,y_R\,d_R

入力例



5 5 8
4 3 6831
1 3 4583
0 0 6592
0 1 3063
3 3 4975
1 3 2049
4 2 2104
2 2 781


解答

# クラスカル法のコード

class UnionFindTree():
    # n要素で初期化
    def __init__(self, n):
        self.n = n
        self.par = list(range(n))
        self.rank = [0] * n

    # 木の根を求める
    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    # xとyの属する集合を併合
    def unite(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return

        if self.rank[x] < self.rank[y]:
            self.par[x] = y
        else:
            self.par[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

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

def kruskal():
    es.sort()    # 辺のコストが小さい順にソートする
    UFT = UnionFindTree(V)    # Union-Findの初期化
    res = 0
    for i in range(E):
        e = es[i]
        if not UFT.same(e[1], e[2]):
            UFT.unite(e[1], e[2])
            res += e[0]
    return res


# ここから蟻本のコード

MAX_R = 50000

# 入力
N, M, R = map(int,input().split())
x = [0] * MAX_R
y = [0] * MAX_R
d = [0] * MAX_R

for i in range(R):
    x[i], y[i], d[i] = map(int,input().split())

V = N + M
E = R
es = []
for i in range(R):
    es.append([-d[i], N + x[i], y[i]])
print(10000 * (N + M) + kruskal())

Pythonで理解する蟻本「2-5 Roadblocks（POJ No.3255）」(p.102)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 Roadblocks（POJ No.3255）」(p.102)
のコードをPythonで書き直したものとなっています。

入力


N:交差点の数\\R:道の本数\\交差点s_iと交差点t_iの間に長さc_iの道が存在する(1\leq{i}\leq{R})


N\,R\\s_1\,t_1\,c_1\\:\\s_R\,t_R\,c_R

入力例



4 4
1 2 100
2 3 250
2 4 200
3 4 100


解答

# 優先度付きキューをインポートする
import heapq

MAX_N = 5000
MAX_R = 100000
INF = float('inf')

# 入力
N, R = map(int,input().split())
G = [[] for _ in range(MAX_N)]    # グラフの隣接リスト表現
for i in range(R):
    s, t, c = map(int,input().split())
    G[s-1].append([t-1, c])
    G[t-1].append([s-1, c])

que = []
dist = [INF] * MAX_N     # 最短経路
dist2 = [INF] * MAX_N    # 二番目の最短経路
dist[0] = 0
heapq.heappush(que, [0, 0])

while len(que) != 0:
    p = heapq.heappop(que)
    d = p[0]
    v = p[1]
    if dist2[v] < d:
        continue
    for i in range(len(G[v])):
        e = G[v][i]
        d2 = d + e[1]
        if dist[e[0]] > d2:
            dist[e[0]], d2 = d2, dist[e[0]]
            heapq.heappush(que, [dist[e[0]], e[0]])
        if dist2[e[0]] > d2 and dist[e[0]] < d2:
            dist2[e[0]] = d2
            heapq.heappush(que, [dist2[e[0]], e[0]])
print(dist2[N - 1])

Pythonで理解する蟻本「2-5 最小全域木問題２（クラスカル法）」(p.101)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 最小全域木問題２（クラスカル法）」(p.101)
のコードをPythonで書き直したものとなっています。

入力


V:頂点数\\E:辺数\\頂点s_iと頂点t_iの間に重みc_iの無向辺が存在する(1\leq{i}\leq{E})


V\,E\\s_1\,t_1\,c_1\\:\\s_E\,t_E\,c_E

入力例

f:id:kuruton456:20201201145758p:plain


7 9
0 1 2
0 4 10
1 2 1
1 3 3
1 5 7
3 5 1
3 6 5
4 5 5
5 6 8


コード

# Union-Find木のコード

class UnionFindTree():
    # n要素で初期化
    def __init__(self, n):
        self.n = n
        self.par = list(range(n))
        self.rank = [0] * n

    # 木の根を求める
    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    # xとyの属する集合を併合
    def unite(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return

        if self.rank[x] < self.rank[y]:
            self.par[x] = y
        else:
            self.par[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

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


# ここから蟻本のコード

V, E = map(int,input().split())    # Vは頂点数、Eは辺数
es = []

for _ in range(E):
    s, t, c = map(int,input().split())
    es.append([c, s, t])    # c(辺のコスト)でソートできるように最初の要素にしておく

def kruskal():
    es.sort()    # 辺のコストが小さい順にソートする
    UFT = UnionFindTree(V)    # Union-Findの初期化
    res = 0
    for i in range(E):
        e = es[i]
        if not UFT.same(e[1], e[2]):
            UFT.unite(e[1], e[2])
            res += e[0]
    return res

res = kruskal()

print(res)

実行結果

17