結果

問題 No.2638 Initial fare
コンテスト
ユーザー katonyonko
提出日時 2024-02-19 21:45:58
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 1,224 ms / 2,000 ms
コード長 1,930 bytes
コンパイル時間 265 ms
コンパイル使用メモリ 81,908 KB
実行使用メモリ 141,552 KB
最終ジャッジ日時 2024-09-29 01:41:26
合計ジャッジ時間 19,571 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 25
権限があれば一括ダウンロードができます

ソースコード

diff # 

import io
import sys
from collections import defaultdict, deque, Counter
from itertools import permutations, combinations, accumulate
from heapq import heappush, heappop
from bisect import bisect_right, bisect_left
from math import gcd
import math

_INPUT = """\
6
5
1 2
2 3
3 4
4 5
8
1 2
2 3
2 4
2 5
5 6
6 7
6 8
2
1 2
"""

def bfs(G,s):
  inf=10**30
  D=[inf]*len(G)
  D[s]=0
  dq=deque()
  dq.append(s)
  while dq:
    x=dq.popleft()
    for y in G[x]:
      if D[y]>D[x]+1:
        D[y]=D[x]+1
        dq.append(y)
  return D

def input():
  return sys.stdin.readline()[:-1]

def solve(test):
  N=int(input())
  G=[[] for _ in range(N)]
  for _ in range(N-1):
    a,b=map(int, input().split())
    a-=1
    b-=1
    G[a].append(b)
    G[b].append(a)
  D=bfs(G,0)
  order=sorted([(D[i],i) for i in range(N)],reverse=True)
  ans=0
  childs=[[0]*3 for _ in range(N)]
  for d,i in order:
    tmp=0
    for c in G[i]:
      if D[c]>d:
        childs[i][1]+=childs[c][0]
        childs[i][2]+=childs[c][1]
        tmp+=1
    childs[i][0]=tmp
    ans+=sum(childs[i])
    ans+=childs[i][0]*(childs[i][0]-1)//2
    for c in G[i]:
      if D[c]>d:
        ans+=childs[i][1]-childs[c][0]
  print(ans)

def random_input():
  from random import randint,shuffle
  N=randint(1,10)
  M=randint(1,N)
  A=list(range(1,M+1))+[randint(1,M) for _ in range(N-M)]
  shuffle(A)
  return (" ".join(map(str, [N,M]))+"\n"+" ".join(map(str, A))+"\n")*3

def simple_solve():
  return []

def main(test):
  if test==0:
    solve(0)
  elif test==1:
    sys.stdin = io.StringIO(_INPUT)
    case_no=int(input())
    for _ in range(case_no):
      solve(0)
  else:
    for i in range(1000):
      sys.stdin = io.StringIO(random_input())
      x=solve(1)
      y=simple_solve()
      if x!=y:
        print(i,x,y)
        print(*[line for line in sys.stdin],sep='')
        break

#0:提出用、1:与えられたテスト用、2:ストレステスト用
main(0)
0