#!/usr/bin/env python3 # 28-Z2-6 Kalamita # Jenda Hadrava # Formát vstupu: # - Na prvním řádku se nachází číslo N určující počet zastávek (mimo # centrální). # - Na druhém řádku je N čísel, i-té číslo udává, na kolikátou zastávku i-tá # navazuje. Číslujeme od 0. Centrální stanici značíme -1. # - Na třetím řádku jsou počty cestujících na jednotlivých stanicích. # Příklad (obrázek ze zadání): # 4 # 1 -1 -1 1 # 2 1 3 5 # Přiklad: # C # /|\ # 1 3 5 # | | | # 7 4 2 # # Vstup: # 6 # -1 -1 -1 0 1 2 # 1 3 5 7 4 2 # Práce s haldou # -------------- # Pokud bychom nechtěli znovuvynalézat kolo a haldu implementovat sami, můžeme # použít knihovnu heapq. Stačí do kódu napsat: # import heapq # A následně používat knihovní verze heappush() a heappop(): # heapq.heappush() a heapq.heappop() # Vložení prvku def heappush(halda, prvek): i = len(halda) halda.append(prvek) # Přidáme prvek na konec pole while (i > 0) and (halda[i] < halda[(i-1) // 2]): # Dokud je třeba, probubláváme nahoru halda[i], halda[(i-1) // 2] = halda[(i-1) // 2], halda[i] i = (i - 1) // 2 # Odebrání minima def heappop(halda): minimum = halda[0] halda[0] = halda[-1] # Vezmeme poslední prvek z pole halda.pop() # Odebere poslední prvek z pole i = 0 while 2 * i + 1 < len(halda): # Probubláváme dolu j = i if halda[2*i + 1] < halda[j]: j = 2*i + 1 if (2*i + 2 < len(halda)) and (halda[2*i + 2] < halda[j]): j = 2*i + 2 # V j máme index nejmenšího z prvků i, 2*i+1 a 2*i+2 if i == j: # Není třeba nic opravovat break halda[i], halda[j] = halda[j], halda[i] # Prohození prvků i = j return minimum N = int(input()) otec_zastavky = list(map(int, input().split())) pocet_cestujicich = list(map(int, input().split())) # Spočítáme syny, abychom mohli snadno poznat, které zastávky jsou listy pocet_synu = [0] * N # Vytvoříme pole nul o N prvcích for i in range(N): if otec_zastavky[i] >= 0: pocet_synu[otec_zastavky[i]] += 1 # Přidáme do haldy všechny listy. Aby byla halda setříděná dle počtu # cestujících, přidáváme do haldy dvojice (počet cestujících, index). halda = [] for i in range(N): if pocet_synu[i] == 0: heappush(halda, (pocet_cestujicich[i], i)) # Poznámka na okraj: Sestavení haldy lze udělat i efektivněji -- v lineárním # čase. V této úloze si tím však již nepomůžeme. Při použití knihovny bychom # volali funkci heapq.heapify(). # Dokud existují nějaké listy, budeme je postupně ubírat a vypisovat while len(halda) > 0: (pocet, i) = heappop(halda) # Vezmeme nejmenší list print(i, " (", pocet, ")", sep="") j = otec_zastavky[i] if j >= 0: # Pokud otec listu není centrální zastávka... pocet_synu[j] -= 1 # U otce snížíme počítadlo synů if pocet_synu[j] == 0: # Odtržením jsme vytvořili nový list heappush(halda, (pocet_cestujicich[j], j)) # Přidáme jej do haldy