Problem:
MathJax.Hub.Config({ tex2jax: {inlineMath: [['$','$'], ['\\(','\\)']]} });

C: ツイート数

問題

AORイカちゃんが利用するSNSであるイカったーでは、投稿のことをツイートと呼ぶ。

そして、イカったーでは、ツイートへの返信が多くなると視認性が悪くなることが懸念されるため、あるツイートが次の規則のいずれかを満たすときに画面にそのツイートを表示する仕様になっている。

規則１. どのツイートへも返信していない

規則２. どのツイートからも返信されていない

規則３. 規則２が適用されたツイートから返信先を順に辿ったとき、 $K$ 回未満で辿り着ける

なお、同じツイートは重複して表示されることはない。

いま、 $N$ 個のツイートがあり、 $A_i$ が $0$ のとき $i$ 番目のツイートは返信でないツイートで、 $A_i$ が $0$ でないとき $i$ 番目のツイートは $A_i$ 番目のツイートへの返信のツイートである。

画面に表示されるツイート数を答えよ。

制約

$1 \le N \le 10^5$

$1 \le K \le 10^5$

$0 \le A_i \lt i (i = 1, 2, \dots, N)$

ツイートは時系列順に与えられる。

入力は全て整数で与えられる。

入力形式

入力は以下の形式で与えられる。

$N \ K$
$A_1$
$A_2$
$\vdots$
$A_N$

出力

画面に表示されるツイート数を出力せよ。また、末尾に改行も出力せよ。

サンプル

サンプル入力 1

6 3
0
0
2
3
4
5

サンプル出力 1

5

この時、表示されるツイートは、図の青いツイートである。 よって、この例の解は $5$ である。

サンプル入力 2

12 2
0
1
0
3
4
3
6
7
0
9
10
11

サンプル出力 2

10

サンプル入力 3

5 10
0
0
0
0
0

サンプル出力 3

5

Solution in Scala (judge verdict: Accepted)

import java.util.Scanner

object Main {
  def solve(sc: => Scanner): Long = {
    val N, K = sc.nextInt
    val A = Array.fill(N)(sc.nextInt - 1)
    val Ans = A.map(X => X == -1)
    val leaf = Array.fill(N)(true)
    for (i <- 0 until N) {
      if (A(i) != -1) (leaf(A(i)) = false)
    }
    val memo = Array.fill(N)(K * 2)
    var i = 0
    while (i < N) {
      if (leaf(i)) {
        var now = i
        var idx = 0
        while (now != -1 && idx < K) {
          Ans(now) = true
          now = A(now)
          idx += 1
          if (now != -1) {
            if (memo(now) <= idx) {
              now = -1
            } else {
              memo(now) = idx
            }
          }
        }
      }
      i += 1
    }

    i = 0
    var ans = 0
    while (i < N) {
      if (Ans(i)) {
        ans += 1
      }
      i += 1
    }
    ans
  }

  def main(args: Array[String]): Unit = {
    val sc: Scanner = new Scanner(System.in)

    println(solve(sc))
  }
}