Submission #3503052

---

Source Code Expand

import strutils
import sequtils
import algorithm
import math
import queues
import tables
import sets
import logging
import future

const INF* = int(1e18 + 373)

proc readLine*(): string = stdin.readLine()
proc readSeq*(): seq[string] = readLine().strip().split()
proc readSeq*(n: Natural): seq[string] =
  result = newSeq[string](n)
  for i in 0..<n:
    result[i] = readLine().strip()
proc readInt1*(): int = readSeq().map(parseInt)[0]
proc readInt2*(): (int, int) =
  let a = readSeq().map(parseInt)
  return (a[0], a[1])
proc readInt3*(): (int, int, int) =
  let a = readSeq().map(parseInt)
  return (a[0], a[1], a[2])
proc readInt4*(): (int, int, int, int) =
  let a = readSeq().map(parseInt)
  return (a[0], a[1], a[2], a[3])
type seq2*[T] = seq[seq[T]]
proc newSeq2*[T](n1, n2: Natural): seq2[T] = newSeqWith(n1, newSeq[T](n2))
type seq3*[T] = seq[seq[seq[T]]]
proc newSeq3*[T](n1, n2, n3: Natural): seq3[T] = newSeqWith(n1, newSeqWith(n2, newSeq[T](n3)))

#------------------------------------------------------------------------------#
type UnionFindTree = object
  p: seq[int]
  h: seq[int]

proc initUnionFindTree(n: int): UnionFindTree =
  result.p = newSeq[int](n)
  result.p.fill(-1)
  result.h = newSeq[int](n)
  result.h.fill(0)

proc find(this: var UnionFindTree, v: int): int =
  if this.p[v] == -1:
    return v
  this.p[v] = find(this, this.p[v])
  return this.p[v]

proc same(this: var UnionFindTree, u, v: int): bool =
  this.find(u) == this.find(v)

proc unite(this: var UnionFindTree, u, v: int) =
  var uRoot = this.find(u)
  var vRoot = this.find(v)
  if uRoot == vRoot:
    return

  if this.h[uRoot] < this.h[vRoot]:
    swap(uRoot, vRoot)

  this.p[vRoot] = uRoot
  if this.h[uRoot] == this.h[vRoot]:
    this.h[uRoot] += 1

#------------------------------------------------------------------------------#

type Edge =
  tuple [ w, u, v: int ]

proc main() =
  let (n, m, k) = readInt3()
  let c = readSeq().map(parseInt).map(it => it - 1)

  var es = newSeq[Edge]()
  for i in 0..<m:
    let (a, b, w) = readInt3()
    es.add((w, a - 1, b - 1))
  es.sort(cmp[Edge])

  var uft = initUnionFindTree(n + k)
  for i in 0..<n:
    if c[i] == -1:
      continue
    uft.unite(i, n + c[i])

  var ans = 0
  var cnt = 0
  for e in es:
    if cnt >= k - 1:
      break

    if uft.same(e.u, e.v):
      continue
    uft.unite(e.u, e.v)
    ans += e.w
    cnt += 1

  if cnt == k - 1:
    echo ans
  else:
    echo -1

main()

Submission Info

Compile Error

Hint: system [Processing]
Hint: Main [Processing]
Hint: strutils [Processing]
Hint: parseutils [Processing]
Hint: sequtils [Processing]
Hint: algorithm [Processing]
Hint: math [Processing]
Hint: times [Processing]
Hint: queues [Processing]
Hint: tables [Processing]
Hint: hashes [Processing]
Hint: etcpriv [Processing]
Hint: sets [Processing]
Hint: os [Processing]
Hint: posix [Processing]
Hint: logging [Processing]
lib/pure/logging.nim(128, 22) Hint: 'Exception' is declared but not used [XDeclaredButNotUsed]
Hint: future [Processing]
Hint: macros [Processing]
Hint:  [Link]
Hint: operation successful (24893 lines compiled; 2.728 sec total; 25.256MB; Release Build) [SuccessX]

Judge Result