<commit_before>object Solution {

  def squareRoot(i: Int, nDigits: Int): BigInt = {
    val limit = BigInt(10).pow(nDigits + 1)
    var a = 5 * BigInt(i)
    var b = BigInt(5)
    while (b < limit) {
      if (a >= b) {
        a -= b
        b += 10
      }
      else {
        a *= 100
        b = (b / 10) * 100 + 5
      }
    }
    b / 100
  }

  def isPerfectRoot(x: Int): Boolean = {
    val sqrt = math.sqrt(x).toInt
    sqrt * sqrt == x
  }

  def main(args: Array[String]) {
    val n = readInt
    val p = readInt
    val sum = (1 to n).filter(!isPerfectRoot(_))
                      .map(x => squareRoot(x, p).toString.map(_.asDigit).sum).sum
    println(sum)
  }
}<commit_msg>Add a BigDecimal version. The first squareRoot function times out for last two cases. The new version times out the last test case. Need another look at the problem.<commit_after>object Solution {
  import java.math.BigDecimal.ROUND_HALF_UP
  import java.math.BigDecimal

  def squareRoot(i: Int, nDigits: Int): BigInt = {
    val limit = BigInt(10).pow(nDigits + 1)
    var a = 5 * BigInt(i)
    var b = BigInt(5)
    while (b < limit) {
      if (a >= b) {
        a = a - b
        b += 10
      }
      else {
        a *= 100
        b = 10 * b - 45
      }
    }
    b / 100
  }

  val ZERO = new BigDecimal("0")
  val TWO = new BigDecimal("2")
  def sqrtBigDecimal(A: BigDecimal, SCALE: Int) = {
    var x0 = ZERO
    var x1 = new BigDecimal(math.sqrt(A.intValue));
    while (!x0.equals(x1)) {
        x0 = x1
        x1 = A.divide(x0, SCALE, ROUND_HALF_UP).add(x0)
        x1 = x1.divide(TWO, SCALE, ROUND_HALF_UP);
    }
    x1
  }

  def isPerfectRoot(x: Int): Boolean = {
    val sqrt = math.sqrt(x).toInt
    sqrt * sqrt == x
  }

  def main(args: Array[String]) {
    val n = readInt
    val p = readInt
    val sum = (1 to n).filter(!isPerfectRoot(_))
      .map(x =>
        sqrtBigDecimal(new BigDecimal(x), p + 1).toString.take(p + 1)
          .foldLeft(0) { (x, y) => if (y == '.') x else x + y.asDigit }).sum
    println(sum)
  }
}
{'fraction_non_alphanumeric': 0.1404109589041096, 'fraction_numerical': 0.026908023483365948, 'mean_word_length': 1.9680696661828738, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4889752', 'n_tokens_mistral': 815, 'n_tokens_neox': 757, 'n_words': 212}