Problem p02660 (python)

Fixed submission:
import numpy as np
import sys
#import copy


def xnxn(n=0):
    """
    args:
        int n: a number of lows to read
    example:
        input   1
        retrun  1

        input   1 2 3
        return  [1,2,3]
        input   1 2 3
                4
        return  [[1,2,3],
                [4]]
    """
    if n == 0:
        temp = list(map(int, input().split()))
        if len(temp) == 1:
            return temp[0]
        elif len(temp) > 1:
            return temp
    else:
        temp = [list(map(int, input().split())) for _ in range(n)]
        return temp


def sp_xnxn(n=0):
    """
    args:
        int n: a number of lows to read
    example
        input   1
        retrun  [1]

        input   123
        return  [1,2,3]

        input   123
                4
        return  [[1,2,3],
                [4]]
    """
    if n == 0:
        return [int(k) for k in xsxs()]
    else:
        return [[int(k) for k in xsxs()] for _ in range(n)]


def xsxs(n=0):
    """
    args:
        int n: a number of lows to read
    example:
        input   AA
        retrun  AA

        input   A BBB CC
        return  ["A","BBB","CC"]

        input   A BB CCC
                D
        return  [["A","BBB","CC"],
                ["D"]]
    """
    if n == 0:
        temp = list(input().split())
        if len(temp) == 1:
            return temp[0]
        elif len(temp) > 1:
            return temp
    else:
        return [list(input().split()) for _ in range(n)]


def sp_xsxs(n=0):
    """
    args:
        int n: a number of lows to read
        example:
            input   A
            retrun  ["A"]

            input   ABC
            return  ["A","B","C"]

            input   ABC
                    D
            return  [["A","B","C"],
                    ["D"]]
    """
    if n == 0:
        return [s for s in xsxs()]
    else:
        return [[s for s in xsxs()] for _ in range(n)]


def p_num(a):

    if a == 1:
        return [[a, 1]]
    else:
        tmp = 2
        ans = []
        while tmp < np.sqrt(a) + 1:
            n = 0
            p = [tmp]

            while a % tmp == 0:
                a = int(a / tmp)
                # p.append(tmp)
                n += 1
                c = 1
            if n != 0:
                p.append(n)
                ans.append(p)
            tmp += 1
            c = 0
        if c == 0 and a != 1:
            ans.append([a, 1])
        return ans


n = xnxn()

if n == 1:
    print(0)
    sys.exit()
else:
    a = p_num(n)
    c = 0
    for num in a:
        i = 1
        while num[1] - i > -1:
            num[1] -= i
            i += 1
            c += 1
    print(c)

Buggy submission (verdict: Runtime Error):
import numpy as np
def p_num_edit(n):
    tmp = n
    ans = []
    for i in range(2, int(np.sqrt(n)) + 1):
        cnt = 0
        while tmp % i == 0:
            cnt += 1
            tmp = int(tmp / i)
        if cnt > 0:
            ans.append([i, cnt])
    if tmp != 1:
        ans.append([tmp, 1])
    if ans == []:
        ans.append([n, 1])
    return ans


n = xnxn()

if n == 1:
    print(0)
    sys.exit()
else:
    a = p_num_edit(n)
    c = 0
    for num in a:
        i = 1
        while num[1] - i > -1:
            num[1] -= i
            i += 1
            c += 1
    print(c)