Pythonで学ぶ蟻本初級 2-4 データ構造 | LI techno 　

Pythonで学ぶ蟻本初級 2-4 データ構造

SHARE

この記事は

けんちょさんのQiita記事「AtCoder 版！蟻本 (初級編)」の2章「2-4 データを工夫して記憶する “データ構造”」にて紹介されている問題をPythonで実行してみたページ

キーワード

  1. priority_queue
  2. 二分探索木
  3. Union-Find木

Pythonによる解答例

例題 2-4-1　Expedition (POJ No.2431)

import heapq
N, L, P = map(int, input().split()) #N:ガソリンスタンドの個数 ,L:距離, P:距離
A=list(map(int, input().split()))
B=list(map(int, input().split()))

#-1倍したB[i]を入れる優先度付きキュー
C=[]
#ans:補給回数, pos:今いる場所, tank:タンクのガソリンの量
ans=0
pos=0
tank=P

for i in range(N):
	d=A[i]-pos #次に進む距離

	#十分な量になるまでガソリンを補給
	while tank

ex1　2017 CODE FESTIVAL THANKS C – Factory

import heapq
N, K = map(int, input().split())

machine=[]
for i in range(N):
	a, b = map(int, input().split())
	heapq.heappush(machine, (a,b))

ans=0
for i in range(K):
	a, b = heapq.heappop(machine)
	ans+=a
	a+=b
	heapq.heappush(machine, (a, b))

print(ans)

ex2　ABC 062 D – 3N Numbers

import heapq
N=int(input())
A=list(map(int, input().split()))
A_rev = [-x for x in A[::-1]] #逆順にして各要素を-1倍する
left, right = A[:N], A_rev[:N]
sum_left_list, sum_right_list = [0]*(N+1), [0]*(N+1)
sum_left, sum_right = sum(left), sum(right)
sum_left_list[0], sum_right_list[0]=sum_left, sum_right
heapq.heapify(left)
heapq.heapify(right)
#left: k=N, N+1,..., 2N-1まで動かして大きい方からN個の和を求める
for k in range(N, 2*N):
	val=A[k]
	heapq.heappush(left, val)
	min_val = heapq.heappop(left)
	sum_left += val
	sum_left -= min_val
	sum_left_list[k-N+1] = sum_left

#right: k=N, N+1,..., 2N-1まで動かして大きい方からN個の和を求める
for k in range(N, 2*N):
	val = A_rev[k]
	heapq.heappush(right, val)
	min_val = heapq.heappop(right)
	sum_right += val
	sum_right -= min_val
	sum_right_list[k-N+1] = sum_right
sum_right_list=sum_right_list[::-1]

ans=-float("inf")
for i in range(N+1):
	ans=max(ans, sum_left_list[i]+sum_right_list[i])
print(ans)

例題 2-4-2　二分探索木 (set, map の練習)

ABC 085 B Kagami Mochi

N=int(input())
D=[int(input()) for _ in range(N)]

se=set([])
for d in D:
	se.add(d)
print(len(se))


#もっと短く書ける
N=int(input())
print(len(set(input() for _ in range(N))))

ABC 091 B Two Colors Card Game

from collections import defaultdict
dd=defaultdict(int)

N=int(input())
for _ in range(N):
  dd[input()]+=1

M=int(input())
for _ in range(M):
  dd[input()]-=1

ans=max(dd.values())
ans=max(ans, 0)
print(ans)

例題 2-4-3　例題 2-4-3　食物連鎖 (POJ No.1182)

ex1　ATC 001 B Union Find

# unionfind
class Uf:
	def __init__(self, N):
		self.p = list(range(N))
		self.rank = [0] * N
		self.size = [1] * N
	#検索 ノード番号を受け取って一番上の親ノードの番号を帰す
	def root(self, x):
		if self.p[x] != x:
			self.p[x] = self.root(self.p[x])

		return self.p[x]
	#同じ集合に属するか判定
	def same(self, x, y):
		return self.root(x) == self.root(y)
	#併合
	def unite(self, x, y):
		# 根を探す
		u = self.root(x)
		v = self.root(y)

		if u == v: return

		#木の高さを比較し、低いほうから高いほうに辺を張る
		if self.rank[u] < self.rank[v]:
			self.p[u] = v
			self.size[v] += self.size[u]
			self.size[u] = 0
		else:
			self.p[v] = u
			self.size[u] += self.size[v]
			self.size[v] = 0
			#木の高さが同じなら片方を1増やす
			if self.rank[u] == self.rank[v]:
				self.rank[u] += 1
	#ノード番号を受け取って、そのノードが含まれている集合のサイズを返す
	def count(self, x):
		return self.size[self.root(x)]

N, Q = map(int, input().split())
uf = Uf(N)
 
for i in range(Q): 
	p, a, b = map(int, input().split())

	if p == 0:
		uf.unite(a, b)

	else:
		print("Yes" if uf.same(a, b) else "No")

ex2　ABC 049 D 連結


# unionfind
class Uf:
	def __init__(self, N):
		self.p = list(range(N))
		self.rank = [0] * N
		self.size = [1] * N
	#検索 ノード番号を受け取って一番上の親ノードの番号を帰す
	def root(self, x):
		if self.p[x] != x:
			self.p[x] = self.root(self.p[x])

		return self.p[x]
	#同じ集合に属するか判定
	def same(self, x, y):
		return self.root(x) == self.root(y)
	#併合
	def unite(self, x, y):
		# 根を探す
		u = self.root(x)
		v = self.root(y)

		if u == v: return

		#木の高さを比較し、低いほうから高いほうに辺を張る
		if self.rank[u] < self.rank[v]:
			self.p[u] = v
			self.size[v] += self.size[u]
			self.size[u] = 0
		else:
			self.p[v] = u
			self.size[u] += self.size[v]
			self.size[v] = 0
			#木の高さが同じなら片方を1増やす
			if self.rank[u] == self.rank[v]:
				self.rank[u] += 1
	#ノード番号を受け取って、そのノードが含まれている集合のサイズを返す
	def count(self, x):
		return self.size[self.root(x)]

from collections import Counter	
N, K, L = map(int, input().split())
uf_load = Uf(N)
uf_rail = Uf(N)

for i in range(K):
	p, q = map(int, input().split())
	p-=1
	q-=1
	uf_load.unite(p, q)
for _ in range(L):
	r, s = map(int, input().split())
	r-=1
	s-=1
	uf_rail.unite(r, s)

union_sets = [(uf_load.root(i), uf_rail.root(i)) for i in range(N)]
counts=Counter(union_sets)
ans = [counts[union_set] for union_set in union_sets]
print(*ans)

ex3　ARC 097 D Equals

# unionfind
class Uf:
	def __init__(self, N):
		self.p = list(range(N))
		self.rank = [0] * N
		self.size = [1] * N
	#検索 ノード番号を受け取って一番上の親ノードの番号を帰す
	def root(self, x):
		if self.p[x] != x:
			self.p[x] = self.root(self.p[x])

		return self.p[x]
	#同じ集合に属するか判定
	def same(self, x, y):
		return self.root(x) == self.root(y)
	#併合
	def unite(self, x, y):
		# 根を探す
		u = self.root(x)
		v = self.root(y)

		if u == v: return

		#木の高さを比較し、低いほうから高いほうに辺を張る
		if self.rank[u] < self.rank[v]:
			self.p[u] = v
			self.size[v] += self.size[u]
			self.size[u] = 0
		else:
			self.p[v] = u
			self.size[u] += self.size[v]
			self.size[v] = 0
			#木の高さが同じなら片方を1増やす
			if self.rank[u] == self.rank[v]:
				self.rank[u] += 1
	#ノード番号を受け取って、そのノードが含まれている集合のサイズを返す
	def count(self, x):
		return self.size[self.root(x)]

N, M = map(int, input().split())
uf=Uf(N)
P = list(map(int, input().split()))
P=[i-1 for i in P]
for i in range(M):
	x, y = map(int, input().split())
	x-=1
	y-=1
	uf.unite(x,y)

#「値iを含む位置jと位置iが同じグループに属している」という条件を満たすiの数を数える
ans=0
for i in range(N):
	if uf.root(P[i]) == uf.root(i):
		ans+=1
print(ans)

ex4　ARC 036 D 偶数メートル

# unionfind
class Uf:
    def __init__(self, N):
        self.p = list(range(N))
        self.rank = [0] * N
        self.size = [1] * N
    #検索 ノード番号を受け取って一番上の親ノードの番号を帰す
    def root(self, x):
        if self.p[x] != x:
            self.p[x] = self.root(self.p[x])

        return self.p[x]
    #同じ集合に属するか判定
    def same(self, x, y):
        return self.root(x) == self.root(y)
    #併合
    def unite(self, x, y):
        # 根を探す
        u = self.root(x)
        v = self.root(y)

        if u == v: return

        #木の高さを比較し、低いほうから高いほうに辺を張る
        if self.rank[u] < self.rank[v]:
            self.p[u] = v
            self.size[v] += self.size[u]
            self.size[u] = 0
        else:
            self.p[v] = u
            self.size[u] += self.size[v]
            self.size[v] = 0
            #木の高さが同じなら片方を1増やす
            if self.rank[u] == self.rank[v]:
                self.rank[u] += 1
    #ノード番号を受け取って、そのノードが含まれている集合のサイズを返す
    def count(self, x):
        return self.size[self.root(x)]




N, Q = map(int, input().split())
uf = Uf(2*N) #i番の頂点からi*2番を偶数用、i*2+1番を奇数用とする

for i in range(Q):
    w, x, y, z = map(int, input().split())
    if w == 1:
        #zが偶数なら、2*xと2*y、2*x+1と2*y+1の点をuniteする
        if not z&1:
            # 青
            uf.unite(2*(x-1), 2*(y-1))
            # 赤
            uf.unite(2*(x-1)+1, 2*(y-1)+1)
        #zが奇数なら、2*xと2*y+1、2*xと2*y+1の点をuniteする
        else:
            uf.unite(2*(x-1), 2*(y-1)+1)
            uf.unite(2*(x-1)+1, 2*(y-1))
    else:
        #2*xと2*yがuniteされているかどうか判定する
        if uf.root(2*(x-1)+1) == uf.root(2*(y-1)+1):
            print('YES')
        else:
            print('NO')

ex5　ABC 087 D People on a Line

class WeighttedUnionFind:
	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 union(self, x, y, w):
		rx = self.find(x)
		ry = self.find(y)

		#xの木の高さ < yの木の高さ
		if self.rank[rx] < self.rank[ry]: self.par[rx] = ry self.weight[rx] = w - self.weight[x] + self.weight[y] #xの木の高さ>=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
	#同じ集合に属するか
	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, M = map(int, input().split())
wuf = WeighttedUnionFind(N)

for _ in range(M):
	l, r, d = map(int, input().split())

	if not wuf.same(l, r):
		wuf.union(l, r, d)

	else:
		if wuf.diff(l, r) !=d:
			print("No")
			exit()
print("Yes")
プログラミングコンテストチャレンジブック第2版

posted with ヨメレバ
秋葉拓哉/岩田陽一 マイナビ出版 2012年01月
Amazon

1 COMMENT

コメントを残す

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です

前の記事

「親友は悪女」10話ネタバレ

次の記事

「親友は悪女」11話　ネタバレ