Code translation with ChatGPT

I wanted to explore the ability of ChatGPT to translate computer code from one language to another. (Or more specifically, my ability to write appropriate prompts for ChatGPT to perform the translation.)

I have a keen interest in integer sequences, so I chose a sequence from the On-Line Encyclopedia of Integer Sequences (OEIS) as a test. Specifically, I chose sequence A000055, which is the number of trees with n unlabeled nodes. The first six terms are 1,1,1,2,3,6, as shown in the figure below.

The OEIS entry has Mathematica code, which is not very helpful to me, because I do not know the language and I do not have access to the software. The Mathematica code is shown below.

s[n_, k_] := s[n, k] = a[n + 1 - k] + If[n < 2k, 0, s[n - k, k]]; a[1] = 1; a[n_] := a[n] = Sum[a[i] s[n-1, i] i, {i, 1, n-1}] / (n-1); Table[a[i] - Sum[a[j] a[i-j], {j, 1, i/2}] + If[OddQ[i], 0, a[i/2] (a[i/2] + 1)/2], {i, 1, 50}] (* Robert A. Russell *)        

Clear as mud, isn't it? I asked ChatGPT to explain the code, but the explanation was at a very low level; it had no big-picture understanding of what the code was doing. But I do know Python pretty well, so I asked ChatGPT to translate the code to Python.

def s(n, k)
    if n < 2*k:
        return a[n+1-k]
    return a[n+1-k] + s(n-k, k)

a = [None] * 51
a[1] = 1

def a(n):
    if a[n] is not None:
        return a[n]
    result = 0
    for i in range(1, n):
        result += a[i] * s(n-1, i) * i
    a[n] = result / (n-1)
    return a[n]

result = []
for i in range(1, 51):
    value = a[i] - sum(a[j] * a[i-j] for j in range(1, i//2+1))
    if i % 2 == 0:
        value += a[i//2] * (a[i//2] + 1) // 2
    result.append(value)
print(result):        

Amazing! Unfortunately, the code doesn't run. Can you see the problem?

ChatGPT is trying to use the same name for the function value a(n) and the array element a[n], and the Python interpreter doesn't understand this. To fix this problem, I asked ChatGPT to rewrite the code without defining a function a(n). Here was the result:

def s(n, k)
    if n < 2*k:
        return a[n+1-k]
    return a[n+1-k] + s(n-k, k)

a = [None] * 51
a[1] = 1

for n in range(2, 51):
    result = 0
    for i in range(1, n):
        result += a[i] * s(n-1, i) * i
    a[n] = result / (n-1)

result = []
for i in range(1, 51):
    value = a[i] - sum(a[j] * a[i-j] for j in range(1, i//2+1))
    if i % 2 == 0:
        value += a[i//2] * (a[i//2] + 1) // 2
    result.append(value)
print(result)        

Now this code actually runs, and it produces the correct answer. It is more verbose than the original Mathematica script, but it is easier to understand. Unfortunately, there is a slight bug. The Python script has a floating-point division, but it should be an integer division. It's hard to blame ChatGPT for this mistake, since it didn't realize that the output should be an a list of integers. So I asked ChatGPT to replace the floating-point division with integer division.

def s(n, k):
??? if n < 2*k:
??????? return a[n+1-k]
??? return a[n+1-k] + s(n-k, k)

a = [None] * 51
a[1] = 1

for n in range(2, 51):
??? result = 0
??? for i in range(1, n):
??????? result += a[i] * s(n-1, i) * i
??? a[n] = result // (n-1)

result = []
for i in range(1, 51):
??? value = a[i] - sum(a[j] * a[i-j] for j in range(1, i//2+1))
??? if i % 2 == 0:
??????? value += a[i//2] * (a[i//2] + 1) // 2
??? result.append(value)
print(result)        

Now the code is correct, but it could be improved. I asked ChatGPT to rewrite the code in idiomatic Python, and to add docstrings and comments. The comments are perhaps overly literal, but they might be helpful for a beginner.

def s(n, k):
??? """
??? Calculate the value of s(n, k).
?? ?
??? Parameters:
??? n (int): The first argument.
??? k (int): The second argument.
?? ?
??? Returns:
??? int: The value of s(n, k).
??? """
??? # If n is less than 2k, return a[n+1-k]
??? if n < 2*k:
??????? return a[n+1-k]
??? # Otherwise, return a[n+1-k] + s(n-k, k)
??? return a[n+1-k] + s(n-k, k)

# Initialize the a list with 51 elements, with the first element set to 1
a = [None] * 51
a[1] = 1

# Calculate the values of the a list
for n in range(2, 51):
??? # Calculate the sum of a[i] * s(n-1, i) * i for all i from 1 to n-1
??? result = sum(a[i] * s(n-1, i) * i for i in range(1, n))
??? # Divide the sum by n-1 and assign it to the nth element of the a list
??? a[n] = result // (n-1)

# Initialize the result list
result = []

# Calculate the values for the result list
for i in range(1, 51):
??? # Calculate the value of a[i] - sum(a[j] * a[i-j] for j in range(1, i//2+1))
??? value = a[i] - sum(a[j] * a[i-j] for j in range(1, i//2+1))
??? # If i is even, add a[i//2] * (a[i//2] + 1) // 2 to the value
??? if i % 2 == 0:
??????? value += a[i//2] * (a[i//2] + 1) // 2
??? # Append the value to the result list
??? result.append(value)

# Print the result list
print(result)        

Overall, I was impressed with the results. ChatGPT succeeded in translating Mathematica code (a niche programming language) into readable Python code, even if it needed some hand-holding.