結果
問題 | No.1002 Twotone |
ユーザー | vwxyz |
提出日時 | 2024-05-03 09:26:34 |
言語 | PyPy3 (7.3.15) |
結果 |
MLE
|
実行時間 | - |
コード長 | 11,787 bytes |
コンパイル時間 | 198 ms |
コンパイル使用メモリ | 82,176 KB |
実行使用メモリ | 1,154,352 KB |
最終ジャッジ日時 | 2024-11-24 08:18:00 |
合計ジャッジ時間 | 124,396 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge2 |
(要ログイン)
テストケース
テストケース表示入力 | 結果 | 実行時間 実行使用メモリ |
---|---|---|
testcase_00 | MLE | - |
testcase_01 | MLE | - |
testcase_02 | MLE | - |
testcase_03 | MLE | - |
testcase_04 | MLE | - |
testcase_05 | MLE | - |
testcase_06 | MLE | - |
testcase_07 | MLE | - |
testcase_08 | MLE | - |
testcase_09 | TLE | - |
testcase_10 | MLE | - |
testcase_11 | MLE | - |
testcase_12 | MLE | - |
testcase_13 | MLE | - |
testcase_14 | AC | 78 ms
68,736 KB |
testcase_15 | AC | 4,913 ms
437,652 KB |
testcase_16 | MLE | - |
testcase_17 | MLE | - |
testcase_18 | AC | 115 ms
78,208 KB |
testcase_19 | TLE | - |
testcase_20 | TLE | - |
testcase_21 | TLE | - |
testcase_22 | AC | 132 ms
78,072 KB |
testcase_23 | TLE | - |
testcase_24 | TLE | - |
testcase_25 | TLE | - |
testcase_26 | AC | 68 ms
65,536 KB |
testcase_27 | AC | 1,275 ms
290,192 KB |
testcase_28 | AC | 2,074 ms
320,824 KB |
testcase_29 | AC | 1,762 ms
323,544 KB |
testcase_30 | AC | 56 ms
57,472 KB |
testcase_31 | AC | 1,750 ms
319,584 KB |
testcase_32 | AC | 2,246 ms
320,044 KB |
testcase_33 | AC | 1,831 ms
320,340 KB |
testcase_34 | MLE | - |
ソースコード
from collections import defaultdict class Graph: def __init__(self,V,edges=None,graph=None,directed=False,weighted=False,inf=float("inf")): self.V=V self.directed=directed self.weighted=weighted self.inf=inf if graph!=None: self.graph=graph """ self.edges=[] for i in range(self.V): if self.weighted: for j,d in self.graph[i]: if self.directed or not self.directed and i<=j: self.edges.append((i,j,d)) else: for j in self.graph[i]: if self.directed or not self.directed and i<=j: self.edges.append((i,j)) """ else: self.edges=edges self.graph=[[] for i in range(self.V)] if weighted: for i,j,d in self.edges: self.graph[i].append((j,d)) if not self.directed: self.graph[j].append((i,d)) else: for i,j in self.edges: self.graph[i].append(j) if not self.directed: self.graph[j].append(i) def Centroid(self,root=0): x=root parents,size=self.SIV_DFS(x,parents=True,subtree_size=True) while True: for y in self.graph[x]: if self.weighted: y,d=y if y==parents[x]: continue if size[y]*2>size[root]: x=y break else: for y in self.graph[x]: if self.weighted: y,d=y if y==parents[x]: continue if size[root]<=2*size[y]: return x,y return x,None def Centroid_Decomposition(self,points=False,edges=False,tree=False,linked_point=False): if points: cd_points=[None]*self.V if edges: cd_edges=[None]*self.V if tree: cd_tree=[]*self.V if linked_point: cd_linked_point=[None]*self.V E=self.edges P=[i for i in range(self.V)] prev_centroid=None stack=[(E,P,None,prev_centroid)] if linked_point else [(E,P,prev_centroid)] while stack: if linked_point: E,P,lp,prev_centroid=stack.pop() else: E,P,prev_centroid=stack.pop() if len(P)==1: centroid=P[0] if edges: cd_edges[centroid]=[] if linked_point: cd_linked_point[centroid]=lp if points: cd_points[centroid]=[centroid] if tree and prev_centroid!=None: cd_tree.append((prev_centroid,centroid)) continue G=Graph(len(P),edges=E,weighted=self.weighted) centroid,_=G.Centroid() if tree and prev_centroid!=None: cd_tree.append((prev_centroid,P[centroid])) parents,tour=G.SIV_DFS(centroid,parents=True,preorder=True) dp=[None]*len(P) EE=[] PP=[] if linked_point: linked_points=[] for i,x in enumerate(G.graph[centroid]): if G.weighted: x,d=x dp[x]=(i,0) EE.append([]) PP.append([P[x]]) if linked_point: linked_points.append(P[x]) for x in tour[1:]: for y in G.graph[x]: if G.weighted: y,d=y if y==parents[x]: continue i,j=dp[x] jj=len(PP[i]) EE[i].append((j,jj,d) if G.weighted else (j,jj)) PP[i].append(P[y]) dp[y]=(i,jj) centroid=P[centroid] if points: cd_points[centroid]=P if edges: cd_edges[centroid]=E if linked_point: cd_linked_point[centroid]=lp if linked_point: for E,P,lp in zip(EE,PP,linked_points): stack.append((E,P,lp,centroid)) else: for E,P in zip(EE,PP): stack.append((E,P,centroid)) retu=() if points: retu+=(cd_points,) if edges: retu+=(cd_edges,) if tree: retu+=(cd_tree,) if linked_point: retu+=(cd_linked_point,) if len(retu)==1: retu=retu[0] return retu def SIV_DFS(self,s,bipartite_graph=False,cycle_detection=False,directed_acyclic=False,euler_tour=False,linked_components=False,lowlink=False,parents=False,postorder=False,preorder=False,subtree_size=False,topological_sort=False,unweighted_dist=False,weighted_dist=False): seen=[False]*self.V finished=[False]*self.V if directed_acyclic or cycle_detection or topological_sort: dag=True if euler_tour: et=[] if linked_components: lc=[] if lowlink: order=[None]*self.V ll=[None]*self.V idx=0 if parents or cycle_detection or lowlink or subtree_size: ps=[None]*self.V if postorder or topological_sort: post=[] if preorder: pre=[] if subtree_size: ss=[1]*self.V if unweighted_dist or bipartite_graph: uwd=[self.inf]*self.V uwd[s]=0 if weighted_dist: wd=[self.inf]*self.V wd[s]=0 stack=[(s,0)] if self.weighted else [s] while stack: if self.weighted: x,d=stack.pop() else: x=stack.pop() if not seen[x]: seen[x]=True stack.append((x,d) if self.weighted else x) if euler_tour: et.append(x) if linked_components: lc.append(x) if lowlink: order[x]=idx ll[x]=idx idx+=1 if preorder: pre.append(x) for y in self.graph[x]: if self.weighted: y,d=y if not seen[y]: stack.append((y,d) if self.weighted else y) if parents or cycle_detection or lowlink or subtree_size: ps[y]=x if unweighted_dist or bipartite_graph: uwd[y]=uwd[x]+1 if weighted_dist: wd[y]=wd[x]+d elif not finished[y]: if (directed_acyclic or cycle_detection or topological_sort) and dag: dag=False if cycle_detection: cd=(y,x) elif not finished[x]: finished[x]=True if euler_tour: et.append(~x) if lowlink: bl=True for y in self.graph[x]: if self.weighted: y,d=y if ps[x]==y and bl: bl=False continue ll[x]=min(ll[x],order[y]) if x!=s: ll[ps[x]]=min(ll[ps[x]],ll[x]) if postorder or topological_sort: post.append(x) if subtree_size: for y in self.graph[x]: if self.weighted: y,d=y if y==ps[x]: continue ss[x]+=ss[y] if bipartite_graph: bg=[[],[]] for tpl in self.edges: x,y=tpl[:2] if self.weighted else tpl if uwd[x]==self.inf or uwd[y]==self.inf: continue if not uwd[x]%2^uwd[y]%2: bg=False break else: for x in range(self.V): if uwd[x]==self.inf: continue bg[uwd[x]%2].append(x) retu=() if bipartite_graph: retu+=(bg,) if cycle_detection: if dag: cd=[] else: y,x=cd cd=self.Route_Restoration(y,x,ps) retu+=(cd,) if directed_acyclic: retu+=(dag,) if euler_tour: retu+=(et,) if linked_components: retu+=(lc,) if lowlink: retu=(ll,) if parents: retu+=(ps,) if postorder: retu+=(post,) if preorder: retu+=(pre,) if subtree_size: retu+=(ss,) if topological_sort: if dag: tp_sort=post[::-1] else: tp_sort=[] retu+=(tp_sort,) if unweighted_dist: retu+=(uwd,) if weighted_dist: retu+=(wd,) if len(retu)==1: retu=retu[0] return retu N,K=map(int,input().split()) edges=[] for i in range(N-1): u,v,c=map(int,input().split()) u-=1;v-=1;c-=1 edges.append((u,v,c)) G=Graph(N,edges=edges,weighted=True) P,E,CD=G.Centroid_Decomposition(points=True,edges=True,tree=True) CD=Graph(N,edges=CD) for c in range(N): if len(P[c])==N: break cd_parents=CD.SIV_DFS(c,parents=True) ans=0 cnt1=[0]*K for x in range(N): memo1=[] child=[y for y in CD.graph[x] if y!=cd_parents[x]] le=len(P[x]) GG=Graph(le,edges=E[x],weighted=True) s=P[x].index(x) parents,tour=GG.SIV_DFS(s,parents=True,preorder=True) dp=[None]*le dp[s]=[] subtree=[None]*le for i in tour: for j,d in GG.graph[i]: if parents[i]==j: continue dp[j]=dp[i][:] if not d in dp[j]: dp[j].append(d) if len(dp[j])>=4: dp[j]=dp[j][:3] if len(dp[j])==2 and dp[j][0]>dp[j][1]: dp[j][0],dp[j][1]=dp[j][1],dp[j][0] if subtree[i]==None: subtree[j]=j else: subtree[j]=subtree[i] subtree_lc=[[] for i in range(le)] for i in range(le): if len(dp[i])<=2 and subtree[i]!=None: subtree_lc[subtree[i]].append(i) cnt2=defaultdict(int) cnt=0 for p in range(le): for i in subtree_lc[p]: if len(dp[i])==2: ans+=cnt2[dp[i][0]*K+dp[i][1]] ans+=cnt2[dp[i][0]] ans+=cnt2[dp[i][1]] elif len(dp[i])==1: ans+=cnt1[dp[i][0]]+cnt-cnt2[(dp[i][0],)] for i in subtree_lc[p]: cnt2[tuple(dp[i])]+=1 if len(dp[i])==2: cnt2[dp[i][0]*K+dp[i][1]]+=1 for j in dp[i]: cnt1[j]+=1 memo1.append(j) elif len(dp[i])==1: cnt2[dp[i][0]]+=1 cnt+=1 for j in memo1: cnt1[j]=0 for i in range(le): if dp[i]!=None and len(dp[i])==2: ans+=1 print(ans)