Hits Different (Codeforces 1829G)
Tags: Math, data structures

In a carnival game, there is a huge pyramid of cans with $$$2023$$$ rows, numbered in a regular pattern as shown. If can $$$9^2$$$ is hit initially, then all cans colored red in the picture above would fall. You throw a ball at the pyramid, and it hits a single can with number $$$n^2$$$. This causes all cans that are stacked on top of this can to fall (that is, can $$$n^2$$$ falls, then the cans directly above $$$n^2$$$ fall, then the cans directly above those cans, and so on). For example, the picture above shows the cans that would fall if can $$$9^2$$$ is hit.What is the sum of the numbers on all cans that fall? Recall that $$$n^2 = n \times n$$$.

Input: The first line contains an integer $$$t$$$ ($$$1 \leq t \leq 1000$$$) — the number of test cases.The only line of each test case contains a single integer $$$n$$$ ($$$1 \leq n \leq 10^6$$$) — it means that the can you hit has label $$$n^2$$$.

Output: For each test case, output a single integer — the sum of the numbers on all cans that fall.Please note, that the answer for some test cases won't fit into 32-bit integer type, so you should use at least 64-bit integer type in your programming language (like long long for C++). For all valid inputs, the answer will always fit into 64-bit integer type.

Note: The first test case is pictured in the statement. The sum of the numbers that fall is $$$$$$1^2 + 2^2 + 3^2 + 5^2 + 6^2 + 9^2 = 1 + 4 + 9 + 25 + 36 + 81 = 156.$$$$$$In the second test case, only the can labeled $$$1^2$$$ falls, so the answer is $$$1^2=1$$$.In the third test case, the cans labeled $$$1^2$$$ and $$$2^2$$$ fall, so the answer is $$$1^2+2^2=1+4=5$$$.In the fourth test case, the cans labeled $$$1^2$$$ and $$$3^2$$$ fall, so the answer is $$$1^2+3^2=1+9=10$$$.In the fifth test case, the cans labeled $$$1^2$$$, $$$2^2$$$, and $$$4^2$$$ fall, so the answer is $$$1^2+2^2+4^2=1+4+16=21$$$.

Sample Inputs:
10
9
1
2
3
4
5
6
10
1434
1000000

Sample Outputs:
156
1
5
10
21
39
46
146
63145186
58116199242129511

Solution (Python):
index = 2
result = [0, 1]
row = 2
while index < 1000000:
    for i in range(index, index + row):
        if i == index:
            summa = result[i - row + 1] + i * i
        elif i == index + row - 1:
            summa = result[index - 1] + i * i
        else:
            summa = result[i - row] + result[i - row + 1] - result[i - 2 * row + 2] + i * i
        result.append(summa)
    index = index + row
    row += 1

def solve():
    n = int(input())
    return result[n]
t = int(input())
results = []
for i in range(t):
    results.append(solve())
for i in results:
    print(i)

Solution (Python):
def main() -> int:
    t: int = int(input())
    pyramid: list = compute_pyramid()
    for _ in range(t):
        solve(pyramid)

def solve(pyramid: list) -> None:
    n: int = int(input())
    print(pyramid[n - 1])

def compute_pyramid() -> list:
    l: list = list()
    n: int = 1000000
    l = [[1]]
    ans: list = list()
    ans.append(1)
    aux: int = 2
    for column in range(1, 1414):
        l.append(list())
        for line in range(column + 1):
            sum: int = 0
            if line == 0:
                sum = aux ** 2 + l[line][column - 1 - line]
            elif column - line == 0:
                sum = aux ** 2 + l[line - 1][column - line]
            else:
                sum = aux ** 2 + l[line][column - 1 - line] + l[line - 1][column - line] - l[line - 1][column - 1 - line]
            ans.append(sum)
            l[line].append(sum)
            aux += 1
    return ans
if __name__ == '__main__':
    main()

Solution (Python):
def main() -> int:
    t: int = int(input())
    pyramid: list = compute_pyramid()
    for _ in range(t):
        solve(pyramid)

def solve(pyramid: list) -> None:
    n: int = int(input())
    print(pyramid[n - 1])

def compute_pyramid() -> list:
    l: list = list()
    values: dict = dict()
    n: int = 1000000
    l = [[1]]
    ans: list = list()
    ans.append(1)
    aux: int = 2
    for column in range(1, 3000):
        l.append(list())
        for line in range(column + 1):
            sum: int = 0
            if line == 0:
                sum = aux ** 2 + l[line][column - 1 - line]
            elif column - line == 0:
                sum = aux ** 2 + l[line - 1][column - line]
            else:
                sum = aux ** 2 + l[line][column - 1 - line] + l[line - 1][column - line] - l[line - 1][column - 1 - line]
            ans.append(sum)
            l[line].append(sum)
            aux += 1
            if aux > n:
                break
        if aux > n:
            break
    return ans
if __name__ == '__main__':
    main()