Count of different numbers divisible by 3 that can be obtained by changing at most one digit

main.py

# '0081' -> 11
# '23' -> 7
# '022' -> 9

def is_div(s: str) -> bool:
  #print(int(s))
  if (int(s) % 3 == 0):
    return True
  return False

# O(n^2)
def solve1(nstr: str) -> int:
  res = 0
  for i in range(len(nstr)):
    #print(int(nstr[i]))
    temp = list(nstr)
    for x in range(10):
      if x == int(nstr[i]):
        continue
      #print(str(x))
      temp[i] = str(x)
      if is_div("".join(temp)):
        res += 1
  if is_div(nstr):
    res += 1
  print("Result: ", res)
  return res

# O(n)
def solve2(nstr: str) -> int:
  result = 0
  sum = 0
  for c in nstr:
    sum += int(c)
  rest = sum % 3
  # Check if initial string is divisible by 3
  if rest == 0:
    result += 1
  # Changing a digit means that sum of all digits
  # may change by minimum of -9 and maximum of 9.
  # We want to add (or subtract) only those numbers,
  # that keep the sum divisible by 3,
  # hence we use (inclusive) range
  # [-9-rest, 9-rest], where step is 3.
  for inc in range(-9-rest, 9-rest+1, 3):
    if inc == 0:
      # We have already taken into account initial string.
      continue
    for i in range(len(nstr)):
      x = int(nstr[i]) + inc
      if x >= 0 and x <= 9:
        result += 1
  print("Result: ", result)
  return result

def main():
  assert(solve2('23') == 7)
  assert(solve2('0081') == 11)
  assert(solve2('022') == 9)
  assert(solve2('235') == 9)

if __name__ == '__main__':
  main()