Mojo π₯ Programming
CPU & GPU algorithm implementations in the Mojo programming language.
Categories
- Arrays & Hashing β 16 problems
- Strings β 13 problems
- Sliding window β 1 problem
- Greedy & Two Pointers β 2 problems
- Search & Sort β 6 problems
- Linked Lists β 1 problem
- Dynamic Programming β 9 problems
- Bit Manipulation & Math β 2 problems
- Grid & Matrix β 3 problems
- Mojo Concepts β 6 problems
- GPU programming β 12 problems
- Utilities β 1 problem
- CUDA β 1 problem
Full source on GitHub
Two Sum Problem
Given an array of integers nums and a target value target,
return the indices of two numbers such that they add up to target.
Approach: Hash Map (One Pass)
The brute-force O(nΒ²) approach checks every pair. We can do better:
- Iterate through the array once.
- For each element, compute
diff = target - nums[i]. - If
diffexists in the hash map, weβve found the pair β return(map[diff], i). - Otherwise, store
nums[i] β iin the map and continue.
This gives O(n) time and O(n) space β optimal for this problem.
A naive O(nΒ²) version (two_sum_costly) is also included for comparison.
def two_sum(nums: List[Int], target: Int) -> Tuple[Int, Int]:
# Default return value: (-1, -1) if no valid pair is found
indices = (-1, -1)
# Early exit: If list has 0 or 1 elements, no pair can be formed
if len(nums) <= 1:
return indices
# Create a dictionary to map each value to its index for quick lookup
# Format: value_indices[value] = index
var value_indices = Dict[Int, Int]() # value -> index
# Iterate through the array
for idx in range(len(nums)):
# Calculate the number needed to reach the target
diff = target - nums[idx]
# If this difference was seen before, we found the pair
if diff in value_indices:
# Retrieve the stored index of the matching number
indices[0] = value_indices.get(diff).value()
# Store the current index as the second of the pair
indices[1] = idx
else:
# Store the current number and its index for future reference
value_indices[nums[idx]] = idx
# Return the result tuple
return indices
def two_sum_costly(nums: List[Int], target: Int) -> Tuple[Int, Int]:
if len(nums) <= 1:
return (-1, -1)
for i in range(len(nums)):
for j in range(i + 1, len(nums)):
if nums[i] + nums[j] == target:
return (i, j)
return (-1, -1)
from std.testing import assert_true, TestSuite
def test_sum() raises:
nums = [2, 7, 11, 15]
target = 9
indices = two_sum(nums, target)
assert_true(indices[0] == 0 and indices[1] == 1, "Assertion failed")
indices = two_sum_costly(nums, target)
assert_true(indices[0] == 0 and indices[1] == 1, "Assertion failed")
target = 18
indices = two_sum(nums, target)
assert_true(indices[0] == 1 and indices[1] == 2, "Assertion failed")
indices = two_sum_costly(nums, target)
assert_true(indices[0] == 1 and indices[1] == 2, "Assertion failed")
target = 100
indices = two_sum(nums, target)
assert_true(indices[0] == -1 and indices[1] == -1, "Assertion failed")
indices = two_sum_costly(nums, target)
assert_true(indices[0] == -1 and indices[1] == -1, "Assertion failed")
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
3SUM
Given an integer array, find all unique triplets (i, j, k) such that nums[i] + nums[j] + nums[k] = 0.
Each triplet is returned with both the values and their original indices.
Approach: Sort + Two Pointers
A naive O(nΒ³) triple loop checks every combination. We can do better:
- Pair each value with its original index, then sort by value.
- Fix one element as the pivot. For each pivot, use two pointers (
left,right) on the remaining subarray. - If the three sum to zero β record the triplet, move both pointers, and skip duplicates.
- If the sum is too low, advance
left; if too high, retreatright. - Early exit: if the pivot value exceeds 0, no later triplet can sum to 0.
This runs in O(nΒ²) time and O(n) space (or O(1) extra ignoring the sort).
A brute-force O(nΒ³) version (find_triplets_bruteforce) is also included for comparison.
from std.random import shuffle
from std.collections import InlineArray
comptime Triplet = InlineArray[Tuple[Int, Int], 3]
# Comparison function used for sorting the array of tuples (value, original_index)
@parameter
def compare_fn(left: Tuple[Int, Int], right: Tuple[Int, Int]) -> Bool:
return (
left[0] < right[0]
) # Compare based on the actual value, not the index
# Function to find all unique triplets that sum up to 0
def find_triplets(nums: List[Int]) -> List[Triplet]:
# Create a new list of tuples (value, original index)
var numbers = List[Tuple[Int, Int]](capacity=len(nums))
for idx in range(len(nums)):
numbers.append((nums[idx], idx)) # Keep original index
# Sort the list based on the value of elements
sort[compare_fn](numbers)
# List to store the result triplets
var result: List[Triplet] = []
# Iterate through the sorted list using a fixed pivot at index `idx`
for idx in range(len(numbers) - 2):
# Early stopping: if current value is greater than 0, we can't find a sum of 0
if numbers[idx][0] > 0:
break
# Skip duplicates: avoid repeated first elements to prevent repeated triplets
if idx > 0 and numbers[idx][0] == numbers[idx - 1][0]:
continue
# Two-pointer approach starts here
var left, right = idx + 1, len(numbers) - 1
while left < right:
# Calculate sum of the triplet
sum = numbers[idx][0] + numbers[left][0] + numbers[right][0]
if sum == 0:
# Found a valid triplet that sums to 0
result.append(
# Triplet(numbers[idx], numbers[left], numbers[right])
[numbers[idx], numbers[left], numbers[right]]
)
# Move both pointers inward
left += 1
right -= 1
# Skip duplicates after finding a valid triplet
while left < right and numbers[left - 1][0] == numbers[left][0]:
left += 1
while (
left < right and numbers[right][0] == numbers[right + 1][0]
):
right -= 1
elif sum < 0:
# Sum is too small, move left pointer to increase it
left += 1
else:
# Sum is too large, move right pointer to decrease it
right -= 1
return result^
# O(nΒ³) brute-force version β checks every (i, j, k) with i < j < k
def find_triplets_bruteforce(nums: List[Int]) -> List[Triplet]:
var result: List[Triplet] = []
for i in range(len(nums) - 2):
for j in range(i + 1, len(nums) - 1):
for k in range(j + 1, len(nums)):
if nums[i] + nums[j] + nums[k] == 0:
result.append([(nums[i], i), (nums[j], j), (nums[k], k)])
return result^
# Helper function to nicely print triplet values and their original indices
def pretty_print(triplets: List[Triplet]):
for triplet in triplets:
print("Elem1 value: ", triplet[0][0], ", index: ", triplet[0][1])
print("Elem2 value: ", triplet[1][0], ", index: ", triplet[1][1])
print("Elem3 value: ", triplet[2][0], ", index: ", triplet[2][1])
print()
# Comparator for sorting Tuple[Int, Int, Int] lexicographically
@parameter
def cmp_tuple3(a: Tuple[Int, Int, Int], b: Tuple[Int, Int, Int]) -> Bool:
if a[0] != b[0]:
return a[0] < b[0]
if a[1] != b[1]:
return a[1] < b[1]
return a[2] < b[2]
# Convert each Triplet to a sorted (a, b, c) value-tuple, discarding indices
def normalize(triplets: List[Triplet]) -> List[Tuple[Int, Int, Int]]:
var result = List[Tuple[Int, Int, Int]]()
for t in triplets:
var a = t[0][0]
var b = t[1][0]
var c = t[2][0]
if a > b:
(a, b) = (b, a)
if b > c:
(b, c) = (c, b)
if a > b:
(a, b) = (b, a)
result.append((a, b, c))
return result^
# Remove adjacent duplicates from a sorted list of value-triplets
def dedup_sorted(
vals: List[Tuple[Int, Int, Int]]
) -> List[Tuple[Int, Int, Int]]:
if len(vals) == 0:
return List[Tuple[Int, Int, Int]]()
var result = List[Tuple[Int, Int, Int]]()
result.append(vals[0])
for i in range(1, len(vals)):
if vals[i] != vals[i - 1]:
result.append(vals[i])
return result^
# Entry point: generates a test list, runs both versions, prints results
def main():
var list = [-3, -3, -2, -1, 0, 1, 2, 2, 3]
shuffle(list)
print("Sorted + Two Pointers:")
var triplets = find_triplets(list)
pretty_print(triplets)
print("Brute-force O(nΒ³):")
var brute = find_triplets_bruteforce(list)
pretty_print(brute)
print("--- Verification ---")
var n1 = normalize(triplets)
var n2 = normalize(brute)
sort[cmp_tuple3](n1)
sort[cmp_tuple3](n2)
n2 = dedup_sorted(n2)
if len(n1) != len(n2):
print("FAIL: length mismatch (", len(n1), "vs", len(n2), ")")
return
for i in range(len(n1)):
if n1[i] != n2[i]:
print("FAIL: mismatch at", i)
return
print("PASS: both versions produce the same", len(n1), "unique triplets")
4SUM
Given an integer array nums and a target value target, return all unique quadruplets
(a, b, c, d) such that a + b + c + d == target and all indices are distinct.
Approach: Sort + Nested Two Pointers
Extends the 3SUM O(nΒ²) approach by adding one more outer loop:
- Sort the array.
- Fix
i, then fixj(> i). For each pair, use two pointers (low,high) on the remaining subarray. - If the four sum to
targetβ record the quadruplet, move both pointers, and skip duplicates. - If the sum is too low, advance
low; if too high, retreathigh. - Skip duplicate values at both the outer loops and the pointer level to avoid repeated results.
Why O(nΒ³)? Two nested loops (i and j) each contribute O(n), and the two-pointer pass on the remainder is O(n) β totaling O(n Β· n Β· n) = O(nΒ³) in the worst case. This is one polynomial degree above 3SUM because k-sum problems generalize to O(n^(k-1)) with this approach.
Why O(log n)? Sorting the array in-place uses O(log n) stack space for the sort implementation (e.g. quicksort recursion). No additional data structures are allocated beyond a few index variables.
comptime Quadruplet = Tuple[Int, Int, Int, Int]
# Function to find all unique quadruplets that sum up to target
def quadruplets(mut nums: List[Int], target: Int) -> List[Quadruplet]:
var result: List[Quadruplet] = []
if len(nums) == 0:
return result^
var length = len(nums)
sort(nums)
for i in range(length - 3):
if i > 0 and nums[i] == nums[i - 1]:
continue
for j in range(i + 1, length - 2):
if j > i + 1 and nums[j - 1] == nums[j]:
continue
var low, high = j + 1, length - 1
while low < high:
four_sum = nums[i] + nums[j] + nums[low] + nums[high]
if four_sum == target:
result.append(
Quadruplet(nums[i], nums[j], nums[low], nums[high])
)
low += 1
high -= 1
while low < high and nums[low] == nums[low - 1]:
low += 1
while low < high and nums[high] == nums[high + 1]:
high -= 1
elif four_sum < target:
low += 1
else:
high -= 1
return result^
def main():
var nums = [1, 0, -1, 0, -2, 2]
var target = 0
var result = quadruplets(nums, target)
for each in result:
print(each[0], each[1], each[2], each[3])
Product of Array Except Self
Given an integer array nums, return a new array where result[i] is the product
of all elements of nums except nums[i].
You must do this without division and in O(n) time.
Approach: Prefix & Suffix Products
The product at index i = (product of everything left of i) Γ (product of everything right of i).
We compute this in two passes with O(1) extra space (excluding the output array):
- Left pass: iterate forward, store the running prefix product at each position.
- Right pass: iterate backward, multiply each position by the running suffix product.
Each pass is O(n), so total time is O(n). Only a few integer variables are used beyond the output array, so space is O(1) auxiliary.
def product_except_self(nums: List[Int]) -> List[Int]:
# If the input list is empty, return it as is
if len(nums) == 0:
return nums.copy()
# Store the length of the input list
var length = len(nums)
# Initialize the result list with all 1s. This will store our final answer.
result = [1] * length
# prefix_product holds the product of all elements to the *left* of the current index
prefix_product = 1
for idx in range(length):
# For each index, store the current prefix product
result[idx] = prefix_product
# Update the prefix product by multiplying it with the current number
prefix_product *= nums[idx]
# suffix_product holds the product of all elements to the *right* of the current index
suffix_product = 1
# Iterate from right to left
for idx in range(length - 1, -1, -1):
# Multiply the result at index with the current suffix product
result[idx] *= suffix_product
# Update the suffix product by multiplying with current number
suffix_product *= nums[idx]
return result^
# Entry point
def main():
# Example input
nums = [1, 2, 3, 4]
# Call the function to get result
result = product_except_self(nums) # Output should be [24, 12, 8, 6]
# Print the result
print(result)
Maximum Subarray Sum (Kadaneβs Algorithm)
Given an integer array nums, find the contiguous subarray (containing at least one element)
with the largest sum, and return that sum.
Approach: Kadaneβs Algorithm (Dynamic Programming)
A brute-force O(nΒ²) check of every subarray is unnecessarily slow. Kadaneβs algorithm solves it in a single pass:
- Maintain
running_sumβ the best sum ending at the current position. - At each step, decide: extend the existing subarray or start fresh from
nums[i]. This isrunning_sum = max(running_sum + nums[i], nums[i]). - Keep
max_sum= the largestrunning_sumseen so far.
The key insight: if running_sum ever drops below the current element alone,
itβs better to start a new subarray from here. This works because a subarray
must be contiguous β you cannot skip elements.
This runs in O(n) time and O(1) space.
# Function to find the subarray with the maximum sum
def max_sum_sub_array(nums: List[Int]) -> Int:
# If the list is empty, return 0 as no subarray exists
if len(nums) == 0:
return 0
# Initialize the running sum and max sum with the first element
# running_sum: current subarray sum being tracked
# max_sum: maximum subarray sum seen so far
var running_sum, max_sum = nums[0], nums[0]
# Iterate over the list starting from the second element
for idx in range(1, len(nums)):
# Decide whether to extend the previous subarray or start a new subarray at current index
running_sum = max(running_sum + nums[idx], nums[idx])
# Update max_sum if the current running_sum is greater
max_sum = max(max_sum, running_sum)
# Return the maximum subarray sum found
return max_sum
from std.testing import assert_true
def main() raises:
nums = [-2, 1, -3, 4, -1, 2, 1, -5, 4]
max_sum = max_sum_sub_array(nums)
assert_true(max_sum == 6, "Assertion failed")
nums = [5, 4, -1, 7, 8]
max_sum = max_sum_sub_array(nums)
assert_true(max_sum == 23, "Assertion failed")
Max Average Subarray
# Find the maximum average of any contiguous subarray of length k(window size) from the array.
def find_max_average(read nums: List[Int], window_size: UInt) -> Float16:
length = len(nums)
if length == 0 or window_size == 0:
return 0.0 # Return 0 if input is invalid
var max_average: Float16 = 0.0
window_sum = 0
# Compute sum of the first 'window_size' elements
for idx in range(window_size):
window_sum += nums[idx]
max_average = Float16(window_sum / window_size) # Initialize max average
# Slide the window over the array
for idx in range(window_size, length):
window_sum += nums[idx] # Add next element
window_sum -= nums[idx - window_size] # Remove the element going out of window
average = Float16(window_sum / window_size) # Current window average
max_average = max(max_average, average) # Update max average if needed
return max_average
def main():
nums = [1, 12, -5, -6, 50, 3]
window_size = 4
max_average = find_max_average(nums, window_size)
debug_assert(max_average == 12.75, "Assertion failed") # Test case check
Max Subarray Product
# Function to find the maximum product of a contiguous subarray
def max_subarray_product(read nums: List[Int]) -> Int:
# Handle edge case: empty array
if len(nums) == 0:
return 0
# Handle edge case: single element
elif len(nums) == 1:
return nums[0]
else:
# Initialize max_product: if first element is 0, set to 1 temporarily
max_product = 1 if nums[0] == 0 else nums[0]
# Track both current max and min products (important for handling negatives)
curr_max, curr_min = 1, 1
# Iterate through all elements
for idx in range(0, len(nums)):
num = nums[idx]
# Reset both max and min when zero is encountered (new subarray starts)
if num == 0:
curr_max, curr_min = 1, 1
continue
# Preserve previous curr_max for updating curr_min
curr_max_copy = curr_max
# Update current max and min by considering:
# - current number alone
# - product of current number with previous max
# - product of current number with previous min (for negatives)
curr_max = max(curr_max * num, curr_min * num, num)
curr_min = min(curr_max_copy * num, curr_min * num, num)
# Update the global max product
max_product = max(max_product, curr_max)
return max_product
# Entry point
def main():
nums = [2, 3, -2, 4] # Expected maximum product subarray: [2, 3] => 6
max_product = max_subarray_product(nums)
print(max_product)
debug_assert(max_product == 6, "Assertion failed")
Given an array of line heights, find the two lines that form the container holding the most water.
def max_area(heights: List[Int]) -> Int:
# If there are fewer than 2 lines, no container can be formed
if len(heights) < 2:
return 0
left, right = 0, len(heights) - 1
max_area = 0
while left < right:
# Height of container is limited by the shorter of the two lines
min_height = min(heights[left], heights[right])
# Calculate area formed between the two lines and update max_area if it's larger
max_area = max(max_area, (right - left) * min_height)
# Move the pointer that's at the shorter line inward to potentially find a taller line
# This can potentially increase the area despite reducing the width
if heights[left] <= heights[right]:
left += 1
else:
right -= 1
return max_area
from std.testing import assert_equal
def main():
heights = [1, 8, 6, 2, 5, 4, 8, 3, 7]
mx_area = max_area(heights)
assert_equal(mx_area, 49, "Assertion failed")
heights = [1, 1]
mx_area = max_area(heights)
assert_equal(mx_area, 1, "Assertion failed")
Largest Number
Arrange non-negative integers to form the largest possible number and return it as a string.
def largest_number(nums: List[Int]) raises -> String:
if len(nums) == 0:
return ""
strs = List[String](capacity=len(nums))
for each in nums:
strs.append(String(each[]))
sort[compare_fn](strs)
result = StringSlice("").join(strs)
return String(Int(result))
@parameter
def compare_fn(left: String, right: String) -> Bool:
return left + right > right + left
from std.testing import assert_true
def main() raises:
nums = [10, 2]
result = largest_number(nums)
assert_true(result == "210", "Assertion failed")
nums = [3, 30, 34, 5, 9]
result = largest_number(nums)
assert_true(result == "9534330", "Assertion failed")
nums = [0, 0, 0, 0, 0]
result = largest_number(nums)
assert_true(result == "0", "Assertion failed")
nums = List[Int]()
result = largest_number(nums)
assert_true(result == "", "Assertion failed")
Remove duplicates from sorted array
This implementation mutates the array in place. Post de-duplication, only the unique elements are retained. After the last unique element, all entries which have been shifted are discarded.
def remove_duplicates(mut nums: List[Int]) -> None:
# If the list has 0 or 1 element, it's already unique
if len(nums) < 2:
return
# `left` points to the position where the next unique element should go
left = 0
# Start from second element and iterate through the list
for right in range(1, len(nums)):
# If a unique value is found (not equal to the previous one),
# move it to the `left + 1` position
if nums[right - 1] != nums[right]:
left += 1
nums[left] = nums[right]
# After all unique elements are placed at the beginning of the list,
# remove all remaining elements beyond the `left` index
for _ in range(len(nums) - 1, left, -1):
_ = nums.pop() # Discard redundant elements
# Import testing helper for assertions
from std.testing import assert_equal
def main() raises:
# Each test validates that the function keeps only unique sorted elements
nums = [1, 1]
remove_duplicates(nums)
assert_equal(nums, [1], "Assertion failed")
nums = [1, 1, 1]
remove_duplicates(nums)
assert_equal(nums, [1], "Assertion failed")
nums = [1, 1, 1, 2]
remove_duplicates(nums)
assert_equal(nums, [1, 2], "Assertion failed")
nums = [1, 1, 1, 2, 3, 3, 3, 5, 5, 6, 6, 8, 8, 8, 9, 10, 10]
remove_duplicates(nums)
assert_equal(nums, [1, 2, 3, 5, 6, 8, 9, 10], "Assertion failed")
nums = [1, 1, 1, 2, 3, 3, 3, 5, 5, 6, 6, 8, 8, 8, 9, 10, 10, 11]
remove_duplicates(nums)
assert_equal(nums, [1, 2, 3, 5, 6, 8, 9, 10, 11], "Assertion failed")
Merge nums2 into nums1 (in-place)
Given two sorted arrays nums1 (size m + n, with m valid elements followed by n zeros) and nums2 (size n), merge them in-place into nums1 as one sorted array.
def merge(mut nums1: List[Int], nums2: List[Int]):
if len(nums1) == 0 or len(nums2) == 0:
return
# Set pointer m to the last valid element in nums1 (i.e., excluding trailing zeros)
m = len(nums1) - len(nums2) - 1
# Set pointer n to the last element of nums2
n = len(nums2) - 1
# Set pointer last to the end of nums1 (i.e., last index where final element will go)
last = len(nums1) - 1
# Traverse both arrays from the end and fill nums1 from the back
while m >= 0 and n >= 0:
if nums1[m] >= nums2[n]:
# If current nums1 element is greater, place it at 'last' and move pointers
nums1[last] = nums1[m]
m -= 1
else:
# Else, place nums2[n] at 'last' and move pointers
nums1[last] = nums2[n]
n -= 1
last -= 1
# If there are leftover elements in nums2 (i.e., nums2 had smaller elements)
while n >= 0:
nums1[last] = nums2[n]
last -= 1
n -= 1
# No need to handle leftover nums1 elements, they are already in place 1
from std.testing import assert_equal
def main() raises:
nums1 = [5, 8, 11, 13, 0, 0, 0]
nums2 = [3, 9, 19]
merge(nums1, nums2)
assert_equal(nums1, [3, 5, 8, 9, 11, 13, 19], "Assertion failed")
nums1 = [1, 2, 3, 0, 0, 0]
nums2 = [2, 5, 6]
merge(nums1, nums2)
assert_equal(nums1, [1, 2, 2, 3, 5, 6], "Assertion failed")
nums1 = [1]
nums2 = []
merge(nums1, nums2)
assert_equal(nums1, [1], "Assertion failed")
Sum 1D Tensor
from layout import Layout, LayoutTensor
from algorithm import vectorize
from sys import simdwidthof
def summer[
type: DType, layout: Layout, //, simdwidth: Int = simdwidthof[type]()
](
tensor: LayoutTensor[type, layout, MutableAnyOrigin],
start: Int = 0,
end: Int = layout.size(),
) -> Scalar[type]:
result = Scalar[type](0)
@parameter
def sum[simd_width: Int](idx: Int):
result += tensor.load[width=simd_width](0, start + idx).reduce_add()
vectorize[sum, simdwidth](end - start)
return result
def main():
from math import iota
comptime elems_count = 1 << 10
var array = InlineArray[Scalar[DType.uint32], elems_count](fill=0)
iota(array.unsafe_ptr(), elems_count)
tensor = LayoutTensor[
DType.uint32, Layout.row_major(1, elems_count), MutableAnyOrigin
](array.unsafe_ptr())
#print(tensor)
start = 1022
end = 1024
#result = summer[16](tensor, start, end)
result = summer(tensor)
print(result)
Evaluate Reverse Polish Notation.
Evaluate the value of an arithmetic expression expressed in Reverse Polish Notation (postfix notation). In RPN the operator follows its operands, eliminating the need for parentheses.
The expression is given as a list of strings, each entry being either an
integer or one of the four operators +, -, *, /.
Algorithm β O(n) time, O(n) space:
- Iterate tokens left to right.
- If the token is an operator, pop two values from the stack (right operand first, then left), apply the operation, and push the result.
- Otherwise the token represents an integer β push it onto the stack.
- After all tokens have been consumed the stack holds exactly one value: the result.
Division truncates toward zero (e.g. -3 / 2 = -1), matching the
convention used by LeetCode problem 150.
Example:
tokens = ["15", "7", "1", "1", "+", "-", "/", "3", "*",
"2", "1", "1", "+", "+", "-"]
Evaluates to 5, equivalent to:
((15 / (7 - (1 + 1))) * 3) - (2 + (1 + 1))
from std.testing import assert_equal, assert_true, TestSuite
# ββ helpers ββββββββββββββββββββββββββββββββββββββββββββββββββββ
def is_operator(token: StaticString) -> Bool:
return token == "+" or token == "-" or token == "*" or token == "/"
def trunc_div(a: Int, b: Int) -> Int:
"""Integer division truncating toward zero.
LeetCode RPN (`/`) requires truncation toward zero (e.g. `-3 / 2 = -1`),
but Mojo's `//` performs floor division, which rounds toward negative
infinity (e.g. `-3 // 2 = -2`).
The correction: when the operands have opposite signs and the division
is inexact, floor rounded one step too far in the negative direction β
adding 1 recovers truncation toward zero.
"""
var q = a // b
if (a < 0) != (b < 0) and q * b != a:
q += 1
return q
# ββ evaluator ββββββββββββββββββββββββββββββββββββββββββββββββββ
def eval_rpn(tokens: List[StaticString]) raises -> Int:
var stack = List[Int](capacity=len(tokens))
for token in tokens:
if is_operator(token):
var right = stack.pop()
var left = stack.pop()
if token == "+":
stack.append(left + right)
elif token == "-":
stack.append(left - right)
elif token == "*":
stack.append(left * right)
else:
stack.append(trunc_div(left, right))
else:
stack.append(Int(token))
return stack[0]
# ββ tests ββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_simple_add() raises:
assert_equal(eval_rpn(["5", "3", "+"]), 8)
def test_simple_sub() raises:
assert_equal(eval_rpn(["10", "4", "-"]), 6)
def test_simple_mul() raises:
assert_equal(eval_rpn(["2", "3", "*"]), 6)
def test_simple_div() raises:
assert_equal(eval_rpn(["8", "4", "/"]), 2)
def test_complex_example() raises:
assert_equal(
eval_rpn([
"15", "7", "1", "1", "+", "-", "/", "3", "*",
"2", "1", "1", "+", "+", "-",
]),
5,
)
def test_multi_ops() raises:
assert_equal(eval_rpn(["2", "1", "+", "3", "*"]), 9)
assert_equal(eval_rpn(["4", "13", "5", "/", "+"]), 6)
def test_negative_result() raises:
assert_equal(eval_rpn(["1", "5", "-"]), -4)
def test_negative_operands() raises:
assert_equal(eval_rpn(["-2", "5", "+"]), 3)
assert_equal(eval_rpn(["-2", "-3", "+"]), -5)
def test_division_trunc_toward_zero() raises:
assert_equal(eval_rpn(["-6", "4", "/"]), -1)
assert_equal(eval_rpn(["6", "-4", "/"]), -1)
assert_equal(eval_rpn(["-6", "-4", "/"]), 1)
assert_equal(eval_rpn(["-7", "2", "/"]), -3)
def test_all_four_ops() raises:
assert_equal(eval_rpn(["2", "3", "+", "5", "*", "4", "-"]), 21)
def test_single_element() raises:
assert_equal(eval_rpn(["42"]), 42)
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Group Anagrams.
Given an array of ASCII strings strs, group the anagrams together.
Two strings are anagrams if one can be rearranged to form the other
(i.e. they share the same character-frequency signature).
Two implementations are provided:
-
group_anagrams_by_sortingβ O(kΒ·n log n) where k is the average string length and n is the number of strings. Sorts each stringβs bytes to produce a canonical key. -
group_anagramsβ O(kΒ·n) counting sort. Builds a 26-element frequency-vector key (assumes lowercase English letters). Faster but restricted to that character set.
Example:
strs = ["eat", "tea", "tan", "ate", "nat", "bat"]
β [["bat"], ["nat", "tan"], ["ate", "eat", "tea"]]
from std.testing import assert_equal, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Approach 1 β sort each string's bytes
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def group_anagrams_by_sorting(strs: List[String]) -> List[List[String]]:
"""Group anagrams using a sorted-bytes key.
Each string is decomposed into its raw bytes, sorted
lexicographically, and the resulting byte list is used as a
dictionary key. All strings sharing the same sorted key belong
to the same anagram group.
"""
if len(strs) == 0:
return List[List[String]]()
var groupings = Dict[List[UInt8], List[String]]()
for s in strs:
# Canonical key: sorted byte sequence of the string.
var key = [byte for byte in s.bytes()]
sort(key)
# Retrieve existing group (or start a new one) and add `s`.
var group = groupings.pop(key, List[String]())
group.append(s)
groupings[key^] = group^
# Collect all groups. Ownership of `groupings` is transferred
# into the list comprehension so each group is moved out once.
return [group.copy() for group in groupings^.values()]
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Approach 2 β frequency-vector key (lowercase only)
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def group_anagrams(strs: List[String]) -> List[List[String]]:
"""Group anagrams using a 26-slot frequency vector.
Each string is reduced to a `List[Int]` of length 26 where slot
`i` holds the count of character `chr(ord('a') + i)`. Strings
with identical frequency vectors are anagrams.
NOTE: only works for lowercase English letters; any other byte
will either silently index out of range or produce a wrong key.
"""
if len(strs) == 0:
return List[List[String]]()
var groupings = Dict[List[Int], List[String]]()
for s in strs:
# Zero-initialised frequency vector, one slot per letter.
var key = List[Int](length=26, fill=0)
for code_point in s.codepoints():
key[Int(code_point) - Int(Codepoint.ord("a"))] += 1
var group = groupings.pop(key, List[String]())
group.append(s)
groupings[key^] = group^
return [group.copy() for group in groupings^.values()]
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# ββ sorting-based approach ββββββββββββββββββββββββββββββββββ
def test_sorting_example() raises:
var result = group_anagrams_by_sorting(
["eat", "tea", "tan", "ate", "nat", "bat"]
)
assert_equal(len(result), 3)
def test_sorting_single() raises:
var result = group_anagrams_by_sorting(["a"])
assert_equal(len(result), 1)
assert_equal(len(result[0]), 1)
def test_sorting_two_groups() raises:
var result = group_anagrams_by_sorting(["a", "b"])
assert_equal(len(result), 2)
def test_sorting_empty_input() raises:
assert_equal(len(group_anagrams_by_sorting(List[String]())), 0)
def test_sorting_empty_string() raises:
var result = group_anagrams_by_sorting([""])
assert_equal(len(result), 1)
assert_equal(len(result[0]), 1)
assert_equal(result[0][0], "")
def test_sorting_all_same() raises:
var result = group_anagrams_by_sorting(["abc", "cab", "bca"])
assert_equal(len(result), 1)
assert_equal(len(result[0]), 3)
# ββ frequency-vector approach βββββββββββββββββββββββββββββββ
def test_freq_example() raises:
var result = group_anagrams(
["eat", "tea", "tan", "ate", "nat", "bat"]
)
assert_equal(len(result), 3)
def test_freq_single() raises:
var result = group_anagrams(["a"])
assert_equal(len(result), 1)
assert_equal(len(result[0]), 1)
def test_freq_two_groups() raises:
var result = group_anagrams(["a", "b"])
assert_equal(len(result), 2)
def test_freq_empty_input() raises:
assert_equal(len(group_anagrams(List[String]())), 0)
def test_freq_empty_string() raises:
var result = group_anagrams([""])
assert_equal(len(result), 1)
assert_equal(len(result[0]), 1)
assert_equal(result[0][0], "")
def test_freq_all_same() raises:
var result = group_anagrams(["abc", "cab", "bca"])
assert_equal(len(result), 1)
assert_equal(len(result[0]), 3)
def test_freq_repeated_strings() raises:
var result = group_anagrams(["ab", "ab", "ba"])
assert_equal(len(result), 1)
assert_equal(len(result[0]), 3)
# ββ cross-verification ββββββββββββββββββββββββββββββββββββββ
def test_both_implementations_agree() raises:
# Test case 1: LeetCode example
var case1 = List[String]()
case1.append(String("eat"))
case1.append(String("tea"))
case1.append(String("tan"))
case1.append(String("ate"))
case1.append(String("nat"))
case1.append(String("bat"))
assert_equal(
len(group_anagrams_by_sorting(case1)),
len(group_anagrams(case1)),
)
# Test case 2: single element
var case2 = List[String]()
case2.append(String("a"))
assert_equal(
len(group_anagrams_by_sorting(case2)),
len(group_anagrams(case2)),
)
# Test case 3: two separate groups
var case3 = List[String]()
case3.append(String("ab"))
case3.append(String("ba"))
case3.append(String("cd"))
assert_equal(
len(group_anagrams_by_sorting(case3)),
len(group_anagrams(case3)),
)
# Test case 4: all same
var case4 = List[String]()
case4.append(String("abc"))
case4.append(String("cab"))
case4.append(String("bca"))
assert_equal(
len(group_anagrams_by_sorting(case4)),
len(group_anagrams(case4)),
)
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Top K Frequent Elements.
Given an integer array nums and an integer k, return the k most
frequent elements. The answer may be returned in any order.
Two implementations are provided:
-
top_k_frequentβ O(n log n) using a max-heap. Count frequencies with aDict, push each (element, count) pair onto a max-heap, then pop the topk. -
top_kβ O(n) using bucket sort. Uses the frequency as an index into a bucket array of lengthlen(nums) + 1(the maximum possible frequency). Iterate buckets from highest frequency downward, collecting elements untilkare gathered.
Example:
nums = [1, 1, 1, 2, 2, 3], k = 2 β [1, 2]
nums = [1], k = 1 β [1]
from std.collections.binary_heap import BinaryHeap
from std.testing import assert_equal, assert_true, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Helper β heap element (frequency-tracking pair)
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
@fieldwise_init
struct Pair(Copyable & Comparable & ImplicitlyDeletable):
"""(element, frequency) pair ordered by frequency for a max-heap.
BinaryHeap is a max-heap β the element with the largest `__lt__`
"wins". We define `__lt__` to compare frequencies so that the
pair with the highest frequency sits at the top.
"""
var pair: Tuple[Int, Int]
def __lt__(self: Self, rhs: Self) -> Bool:
return self.pair[1] < rhs.pair[1]
def __eq__(self: Self, other: Self) -> Bool:
return self.pair[1] == other.pair[1]
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Approach 1 β max-heap (O(n log n))
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def top_k_frequent(nums: List[Int], k: Int) -> List[Int]:
"""Top `k` frequent elements via a max-heap.
1. Build a frequency dict.
2. Push every (element, frequency) pair onto a max-heap.
3. Pop the top `k` elements β these are the k most frequent.
"""
if len(nums) == 0 or not k > 0 or len(nums) < k:
return []
var freqs = Dict[Int, Int]()
for num in nums:
freqs[num] = 1 + freqs.get(num, 0)
if len(freqs) < k:
return []
var heap = BinaryHeap[Pair]()
# Transfer ownership of freqs to avoid retaining the dict.
for item in freqs^.items():
heap.push(Pair((item.key, item.value)))
var result = List[Int](capacity=k)
for _ in range(k):
# Pop returns the pair with highest frequency.
result.append(heap.pop().pair[0])
return result^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Approach 2 β bucket sort (O(n))
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def top_k(nums: List[Int], k: Int) -> List[Int]:
"""Top `k` frequent elements via bucket sort.
1. Build a frequency dict.
2. Use the frequency as an index into a bucket array (`bucket`).
A number appearing `f` times is placed in bucket `f`.
3. Walk buckets from highest frequency downward, collecting elements
until `k` are gathered.
"""
if len(nums) == 0 or not k > 0 or len(nums) < k:
return []
var freqs = Dict[Int, Int]()
for num in nums:
freqs[num] = 1 + freqs.get(num, 0)
if len(freqs) < k:
return []
# Bucket array: index = frequency, value = list of elements with that freq.
# NOTE: `List[List[Int]]` with `fill` shares the inner list across all slots,
# so we must build each bucket independently to avoid cross-bucket aliasing.
var bucket = List[List[Int]](capacity=len(nums) + 1)
for _ in range(len(nums) + 1):
bucket.append(List[Int]())
for item in freqs.items():
var num = item.key
var index = item.value
bucket[index].append(num)
var result = List[Int](capacity=k)
# Walk from highest possible frequency down to 1.
for right in range(len(bucket) - 1, 0, -1):
for n in bucket[right]:
result.append(n)
if len(result) == k:
return result^
return result^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# ββ max-heap approach βββββββββββββββββββββββββββββββββββββββ
def test_heap_example_1() raises:
assert_equal(top_k_frequent([1, 1, 1, 2, 2, 3], 2), [1, 2])
def test_heap_example_2() raises:
assert_equal(top_k_frequent([1], 1), [1])
def test_heap_example_3() raises:
var result = top_k_frequent([1, 2, 1, 2, 1, 2, 3, 1, 3, 2], 2)
assert_equal(len(result), 2)
assert_true(1 in result)
assert_true(2 in result)
def test_heap_k_larger_than_unique() raises:
assert_equal(top_k_frequent([1, 1, 2], 5), [])
def test_heap_empty_input() raises:
assert_equal(top_k_frequent([], 1), [])
def test_heap_all_same_frequency() raises:
var result = top_k_frequent([1, 2, 3, 4], 2)
assert_equal(len(result), 2)
def test_heap_single_repeated() raises:
assert_equal(top_k_frequent([5, 5, 5, 5], 1), [5])
def test_heap_k_is_zero() raises:
assert_equal(top_k_frequent([1, 2, 3], 0), [])
def test_heap_negative_k() raises:
assert_equal(top_k_frequent([1, 2, 3], -1), [])
def test_heap_negative_numbers() raises:
var result = top_k_frequent([-1, -1, -2, -2, -2, -3], 2)
assert_equal(len(result), 2)
assert_true(-2 in result)
# ββ bucket-sort approach ββββββββββββββββββββββββββββββββββββ
def test_bucket_example_1() raises:
assert_equal(top_k([1, 1, 1, 2, 2, 3], 2), [1, 2])
def test_bucket_example_2() raises:
assert_equal(top_k([1], 1), [1])
def test_bucket_example_3() raises:
var result = top_k([1, 2, 1, 2, 1, 2, 3, 1, 3, 2], 2)
assert_equal(len(result), 2)
assert_true(1 in result)
assert_true(2 in result)
def test_bucket_k_larger_than_unique() raises:
assert_equal(top_k([1, 1, 2], 5), [])
def test_bucket_empty_input() raises:
assert_equal(top_k([], 1), [])
def test_bucket_all_same_frequency() raises:
var result = top_k([1, 2, 3, 4], 2)
assert_equal(len(result), 2)
def test_bucket_single_repeated() raises:
assert_equal(top_k([5, 5, 5, 5], 1), [5])
def test_bucket_k_is_zero() raises:
assert_equal(top_k([1, 2, 3], 0), [])
def test_bucket_negative_k() raises:
assert_equal(top_k([1, 2, 3], -1), [])
def test_bucket_negative_numbers() raises:
var result = top_k([-1, -1, -2, -2, -2, -3], 2)
assert_equal(len(result), 2)
assert_true(-2 in result)
# ββ cross-verification ββββββββββββββββββββββββββββββββββββββ
def _top_k_freq(nums: List[Int], k: Int, idx: Int) -> Int:
"""Compute the k-th highest frequency in `nums`.
Used to verify that an implementation's answer consists solely of
elements whose frequency is at least this threshold.
"""
var freqs = Dict[Int, Int]()
for num in nums:
freqs[num] = 1 + freqs.get(num, 0)
var freq_list = List[Int](capacity=len(freqs))
for item in freqs.items():
freq_list.append(item.value)
sort(freq_list) # ascending
return freq_list[len(freq_list) - k]
def _valid_top_k_result(nums: List[Int], k: Int, result: List[Int]) -> Bool:
"""Check that `result` is a valid answer for the top-k problem.
Every element in `result` must have frequency >= the k-th highest
frequency in `nums`. This accommodates ties where different
implementations may pick different elements.
"""
if len(result) != k:
return False
var threshold = _top_k_freq(nums, k, 0)
var freqs = Dict[Int, Int]()
for num in nums:
freqs[num] = 1 + freqs.get(num, 0)
for x in result:
if freqs.get(x, 0) < threshold:
return False
return True
def test_heap_only_valid_results() raises:
var cases = List[List[Int]]()
var ks = List[Int]()
cases.append([1, 1, 1, 2, 2, 3]); ks.append(2)
cases.append([1]); ks.append(1)
cases.append([1, 2, 1, 2, 1, 2, 3, 1, 3, 2]); ks.append(2)
cases.append([5, 5, 5, 5]); ks.append(1)
cases.append([1, 2, 3, 4]); ks.append(2)
cases.append([-1, -1, -2, -2, -2, -3]); ks.append(2)
cases.append([1, 1, 2, 2, 3, 3, 3]); ks.append(2)
for i in range(len(cases)):
assert_true(
_valid_top_k_result(cases[i], ks[i], top_k_frequent(cases[i], ks[i]))
)
def test_bucket_only_valid_results() raises:
var cases = List[List[Int]]()
var ks = List[Int]()
cases.append([1, 1, 1, 2, 2, 3]); ks.append(2)
cases.append([1]); ks.append(1)
cases.append([1, 2, 1, 2, 1, 2, 3, 1, 3, 2]); ks.append(2)
cases.append([5, 5, 5, 5]); ks.append(1)
cases.append([1, 2, 3, 4]); ks.append(2)
cases.append([-1, -1, -2, -2, -2, -3]); ks.append(2)
cases.append([1, 1, 2, 2, 3, 3, 3]); ks.append(2)
for i in range(len(cases)):
assert_true(
_valid_top_k_result(cases[i], ks[i], top_k(cases[i], ks[i]))
)
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Longest Consecutive Sequence.
Given an unsorted array of integers nums, return the longest contiguous
sequence that can be formed from its elements (the sequence itself, not
just its length). Runs in O(n) time.
Algorithm β O(n) time, O(n) space:
- Insert all numbers into a
Setfor O(1) membership checks. - For each number, check if
n - 1exists in the set. If it does,nis not the start of a sequence β skip it. - If
nis a sequence start, walk upward (n+1,n+2, β¦) while consecutive numbers exist in the set, building the sequence. - Keep the longest sequence seen.
This works because each number belongs to exactly one contiguous block, and only the smallest element of each block triggers the walk, so every element is examined at most twice (once in the outer loop, once during a walk) β O(n) total.
Example:
nums = [100, 4, 200, 1, 3, 2] β [1, 2, 3, 4]
nums = [0, 3, 7, 2, 5, 8, 4, 6, 0, 1] β [0, 1, 2, 3, 4, 5, 6, 7, 8]
from std.collections import Set
from std.testing import assert_equal, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Implementation
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def longest_consecutive_seq(nums: List[Int]) -> List[Int]:
"""Return the longest consecutive sequence contained in `nums`.
Uses a hash set for O(1) lookups and only initiates a walk from
numbers that are the start of a sequence (no predecessor in the
set), guaranteeing O(n) overall time.
"""
if len(nums) == 0 or len(nums) == 1:
return nums.copy()
var uniques = Set(nums)
var longest: List[Int] = []
for n in nums:
# Only start a new sequence if n-1 is absent β otherwise n is
# part of a block already being handled from a smaller start.
if n - 1 not in uniques:
# If the current longest sequence already starts at `n`,
# this block has already been explored β skip redundant work.
if longest and longest[0] == n:
continue
var curr = [n]
var next = n + 1
while next in uniques:
curr.append(next)
next += 1
longest = curr^ if len(curr) > len(longest) else longest^
return longest^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_example_1() raises:
assert_equal(
longest_consecutive_seq([100, 4, 200, 1, 3, 2]),
[1, 2, 3, 4],
)
def test_example_2_with_duplicates() raises:
assert_equal(
longest_consecutive_seq([0, 3, 7, 2, 5, 8, 4, 6, 0, 1]),
[0, 1, 2, 3, 4, 5, 6, 7, 8],
)
def test_example_3() raises:
assert_equal(
longest_consecutive_seq([1, 0, 1, 2]),
[0, 1, 2],
)
def test_empty_input() raises:
assert_equal(longest_consecutive_seq([]), [])
def test_single_element() raises:
assert_equal(longest_consecutive_seq([42]), [42])
def test_two_consecutive() raises:
assert_equal(longest_consecutive_seq([1, 2]), [1, 2])
def test_two_non_consecutive() raises:
var result = longest_consecutive_seq([1, 3])
assert_equal(len(result), 1)
def test_all_duplicates() raises:
assert_equal(longest_consecutive_seq([5, 5, 5, 5]), [5])
def test_negative_numbers() raises:
assert_equal(
longest_consecutive_seq([-5, -4, -3, 0, 1]),
[-5, -4, -3],
)
def test_mixed_negatives_and_positives() raises:
assert_equal(
longest_consecutive_seq([-1, 0, 1, 5, 6, 7, 8]),
[5, 6, 7, 8],
)
def test_long_range_0_to_9() raises:
assert_equal(
longest_consecutive_seq([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
)
def test_unsorted_input() raises:
assert_equal(
longest_consecutive_seq([9, 3, 5, 4, 6, 7, 8, 2, 1, 0]),
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
)
def test_multiple_blocks() raises:
var result = longest_consecutive_seq([10, 11, 12, 20, 21, 30])
assert_equal(len(result), 3)
def test_single_is_also_longest_with_multi_element_blocks() raises:
var result = longest_consecutive_seq([1, 10, 11, 20])
assert_equal(len(result), 2)
assert_equal(result, [10, 11])
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
String to Integer (atoi)
Implement the atoi(string s) function, which converts a string to a signed integer.
def atoi(s: String) raises -> Int:
# Return 0 for an empty string (edge case)
if len(s) == 0:
return 0
buffer = String() # Buffer to collect valid characters forming the number
digits = String("0123456789") # Valid digit characters
idx = 0 # Index for scanning the string
prelude = True # Indicates we're in the whitespace/prefix-skipping phase
while idx < len(s):
# Skip leading whitespaces
while prelude and idx < len(s) and (s[idx] == " "):
idx += 1
continue
# If we have skipped whitespaces, check for sign or digit
if prelude and idx < len(s):
# If current char is '+' or '-' or a digit, add it to buffer
if s[idx] == "-" or s[idx] == "+" or s[idx] in digits:
if not s[idx] == "+":
buffer.__iadd__(s[idx])
idx += 1
prelude = False # Exit prelude phase after processing sign or first digit
continue
else:
# Invalid character before number starts; exit parsing
break
# If we encounter a non-digit (excluding whitespace in middle), break
if s[idx] != " " and s[idx] not in digits:
break
# If it's a digit, add to buffer
if s[idx] in digits:
buffer.__iadd__(s[idx])
idx += 1
# Now buffer contains something like "-123", "456", "+789", etc.
number, factor = 0, 1
# Convert from left to right (excluding the first char which could be a sign)
for idx in range(len(buffer) - 1, 0, -1):
number = number + Int(buffer[idx]) * factor
factor *= 10
# Handle the first character (either a sign or a digit)
number = (
-1 * number if buffer[0] == "-" else number + Int(buffer[0]) * factor
) if len(buffer) > 1 else number
return number
from std.testing import assert_equal
def main() raises:
s = " - 13 37 c0d3"
number = atoi(s)
assert_equal(number, -1337)
s1 = "13 37 c0d3"
number = atoi(s1)
assert_equal(number, 1337)
s2 = "1337c0d3"
number = atoi(s2)
assert_equal(number, 1337)
s3 = " -042"
number = atoi(s3)
assert_equal(number, -42)
s4 = "42"
number = atoi(s4)
assert_equal(number, 42)
s5 = "0-1"
number = atoi(s5)
assert_equal(number, 0)
s6 = "words and 987"
number = atoi(s6)
assert_equal(number, 0)
s7 = " words and 987"
number = atoi(s7)
assert_equal(number, 0)
s8 = "+987"
number = atoi(s8)
assert_equal(number, 987)
s9 = " + 98700 www"
number = atoi(s9)
assert_equal(number, 98700)
s10 = " - 98700 www"
number = atoi(s10)
assert_equal(number, -98700)
Longest Substring Without Repeating Characters.
Given an ASCII string s, find the length of the longest contiguous
substring that contains no duplicate characters.
Uses a sliding-window approach with a Set to track characters in
the current window (O(n) time, O(min(n, |Ξ£|)) space).
Example:
s = "abcabcbb" β 3 ("abc")
s = "bbbbb" β 1 ("b")
s = "pwwkew" β 3 ("wke")
from std.collections import Set
from std.testing import assert_equal, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Sliding-window implementation
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def len_longest_substr_no_char_repeats(ascii_s: String) raises -> Int:
"""Longest substring length with all unique characters (sliding window).
Maintains a Set `seen` over the current window `[left, idx]`.
When a duplicate character is encountered, shrink the window from
the left until the duplicate is removed.
"""
var s = ascii_s.as_bytes()
var n = len(s)
if n == 0 or n == 1:
# Early return without going thru ceremonies
return n
# Seed the window with the first character.
var seen = Set(s[0])
var left = 0
var max_length = 1
for idx in range(1, n):
# Shrink window from the left until the duplicate is gone.
while s[idx] in seen:
seen.remove(s[left])
left += 1
seen.add(s[idx])
# Set size equals current window length (no duplicates).
max_length = max(max_length, len(seen))
return max_length
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_example_cases() raises:
assert_equal(len_longest_substr_no_char_repeats("abcabcbb"), 3)
assert_equal(len_longest_substr_no_char_repeats("bbbbb"), 1)
assert_equal(len_longest_substr_no_char_repeats("pwwkew"), 3)
def test_word_no_repeats() raises:
assert_equal(len_longest_substr_no_char_repeats("current"), 4)
def test_empty_string() raises:
assert_equal(len_longest_substr_no_char_repeats(""), 0)
def test_single_char() raises:
assert_equal(len_longest_substr_no_char_repeats("a"), 1)
def test_two_unique_chars() raises:
assert_equal(len_longest_substr_no_char_repeats("ab"), 2)
def test_two_same_chars() raises:
assert_equal(len_longest_substr_no_char_repeats("aa"), 1)
def test_all_unique() raises:
assert_equal(len_longest_substr_no_char_repeats("abcdef"), 6)
def test_repeat_at_end() raises:
assert_equal(len_longest_substr_no_char_repeats("abca"), 3)
def test_repeat_in_middle() raises:
assert_equal(len_longest_substr_no_char_repeats("abacdef"), 6)
def test_whole_string_is_answer() raises:
assert_equal(len_longest_substr_no_char_repeats("abcdbef"), 5)
def test_special_chars() raises:
assert_equal(len_longest_substr_no_char_repeats(" !@#$%"), 6)
assert_equal(len_longest_substr_no_char_repeats("a b c"), 3)
def test_unicode_not_tested() raises:
# Operates on raw bytes β multi-byte UTF-8 chars like 'Γ©' (2 bytes)
# are treated as distinct bytes, not a single character.
assert_equal(len_longest_substr_no_char_repeats("cafΓ©"), 5)
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Longest Unique-Character Subarray.
Given an ASCII string s, return the longest contiguous substring that
contains no repeated characters.
Uses a sliding-window with a Dict mapping each character to its most
recent index (O(n) time, O(min(n, |Ξ£|)) space).
This is the βreturn the substringβ variant of the classic problem β
compare with longest_substr_no_char_repeats which returns only the
length.
Example:
s = "12345678911" β "123456789"
s = "abcabcbb" β "abc"
from std.testing import assert_equal, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Sliding-window implementation
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def longest_unique_char_subarray(s: String) raises -> String:
"""Longest substring with all unique characters (returns substring).
Expands the window one character at a time (`right` pointer).
When a duplicate is found, the `left` pointer jumps past the
previous occurrence, maintaining a duplicate-free window.
Tracks the start and length of the longest window seen.
"""
var bytes = s.as_bytes()
var n = len(bytes)
if n == 0 or n == 1:
return s
# Map each byte to its most recent index in the current window.
var seen = Dict[UInt8, Int]()
var left = 0
var max_start = 0
var max_length = 0
for right, char in enumerate(bytes):
# If char was seen inside the current window, jump past it.
if char in seen and seen[char] >= left:
left = seen[char] + 1
# Record (or update) the position of this character.
seen[char] = right
var curr_length = right - left + 1
if curr_length > max_length:
max_length = curr_length
max_start = left
return String(from_utf8=bytes[max_start : max_start + max_length])
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_example_with_numbers() raises:
assert_equal(longest_unique_char_subarray("12345678911"), "123456789")
def test_example_abcabcbb() raises:
assert_equal(longest_unique_char_subarray("abcabcbb"), "abc")
def test_example_bbbbb() raises:
assert_equal(longest_unique_char_subarray("bbbbb"), "b")
def test_example_pwwkew() raises:
assert_equal(longest_unique_char_subarray("pwwkew"), "wke")
def test_empty_string() raises:
assert_equal(longest_unique_char_subarray(""), "")
def test_single_char() raises:
assert_equal(longest_unique_char_subarray("x"), "x")
def test_two_unique() raises:
assert_equal(longest_unique_char_subarray("ab"), "ab")
def test_two_same() raises:
assert_equal(longest_unique_char_subarray("aa"), "a")
def test_all_unique() raises:
assert_equal(longest_unique_char_subarray("abcdef"), "abcdef")
def test_answer_in_middle() raises:
assert_equal(longest_unique_char_subarray("abca"), "abc")
def test_answer_at_end() raises:
assert_equal(longest_unique_char_subarray("aabc"), "abc")
def test_longer_substring_later() raises:
assert_equal(longest_unique_char_subarray("abacdef"), "bacdef")
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Generate All Substrings.
Given an ASCII string, return every possible contiguous substring. A substring is defined by a starting and ending index and is non-empty.
Three implementations are provided, each O(nΒ²) time and O(nΒ²) space (since there are n(n+1)/2 substrings total):
gen_all_sub_strs_loopβ nested forβloop (most explicit).gen_all_sub_strs_list_comprehensionβ Mojo list comprehension.gen_all_sub_strs_by_lengthβ shortest substrings first.
Example:
gen_all_sub_strs_loop("abc")
β ["a", "ab", "abc", "b", "bc", "c"]
from std.testing import assert_equal, TestSuite
from std.collections import Set
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Implementations
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def gen_all_sub_strs_loop(s: String) raises -> List[String]:
"""Generate all substrings via explicit nested loops (O(nΒ²))."""
var bytes = s.as_bytes()
var n = len(bytes)
# Pre-allocate capacity: n(n+1)/2 is the exact number of substrings.
var substrings = List[String](capacity=((n * (n + 1)) // 2))
for start in range(n):
for end in range(start + 1, n + 1):
# Span slice [start, end) is exclusive on the right.
substrings.append(String(from_utf8=bytes[start:end]))
return substrings^
def gen_all_sub_strs_list_comprehension(s: String) raises -> List[String]:
"""Generate all substrings using a Mojo list comprehension."""
var bytes = s.as_bytes()
var n = len(bytes)
return [
String(from_utf8=bytes[start:end])
for start in range(n)
for end in range(start + 1, n + 1)
]
def gen_all_sub_strs_by_length(s: String) raises -> List[String]:
"""Generate all substrings ordered by increasing length.
Unlike the other two, this emits single-character substrings first,
then pairs, triples, etc. The order is useful when searching for
the shortest substring matching some predicate (e.g. minimum window).
"""
var n = s.byte_length()
var substrings = List[String](capacity=((n * (n + 1)) // 2))
for length in range(1, n + 1):
# For a fixed length, every possible start offset
# produces one substring of that exact length.
for start in range(n - length + 1):
var end = start + length
substrings.append(String(s[byte=start:end]))
return substrings^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Helpers
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def string_list(items: List[StringSlice[StaticConstantOrigin]]) -> List[String]:
"""Convert literal string slices to a `List[String]` for testing."""
var result = List[String](capacity=len(items))
for item in items:
result.append(String(item))
return result^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_empty_string() raises:
assert_equal(len(gen_all_sub_strs_loop("")), 0)
assert_equal(len(gen_all_sub_strs_list_comprehension("")), 0)
assert_equal(len(gen_all_sub_strs_by_length("")), 0)
def test_single_char() raises:
var e = string_list(["a"])
assert_equal(gen_all_sub_strs_loop("a"), e)
assert_equal(gen_all_sub_strs_list_comprehension("a"), e)
assert_equal(gen_all_sub_strs_by_length("a"), e)
def test_two_chars() raises:
var e = string_list(["a", "ab", "b"])
assert_equal(gen_all_sub_strs_loop("ab"), e)
assert_equal(gen_all_sub_strs_list_comprehension("ab"), e)
# gen_all_sub_strs_by_length uses a different order (tested separately)
def test_three_chars() raises:
var e = string_list(["a", "ab", "abc", "b", "bc", "c"])
assert_equal(gen_all_sub_strs_loop("abc"), e)
assert_equal(gen_all_sub_strs_list_comprehension("abc"), e)
def test_all_implementations_agree() raises:
var inputs = string_list(["", "a", "ab", "xyz", "hello", "abcde"])
for s in inputs:
var loop_result = gen_all_sub_strs_loop(s)
var comp_result = gen_all_sub_strs_list_comprehension(s)
var by_len_result = gen_all_sub_strs_by_length(s)
# Loop and comprehension should produce identical order.
assert_equal(loop_result, comp_result)
# By-length may differ in order, but must contain the same set.
assert_equal(
Set(by_len_result),
Set(loop_result),
)
def test_by_length_order() raises:
var e = string_list([
"a", "b", "c", "d", # length 1
"ab", "bc", "cd", # length 2
"abc", "bcd", # length 3
"abcd", # length 4
])
assert_equal(gen_all_sub_strs_by_length("abcd"), e)
def test_substring_count() raises:
# A string of length n has exactly n(n+1)/2 substrings.
for n in range(1, 7):
var s = String("x" * n)
var expected_count = (n * (n + 1)) // 2
assert_equal(len(gen_all_sub_strs_loop(s)), expected_count)
assert_equal(len(gen_all_sub_strs_list_comprehension(s)), expected_count)
assert_equal(len(gen_all_sub_strs_by_length(s)), expected_count)
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Longest Palindromic Substring
Given a string s, return the longest palindromic substring therein
def longest_palindrome(s: String) -> String:
if len(s) == 0:
return s
longest = String("")
for i in range(len(s)):
left, right = i, i
find_longest(left, right, s, longest)
left, right = i, i + 1
find_longest(left, right, s, longest)
return longest
def find_longest(mut left: Int, mut right: Int, s: String, mut longest: String):
while 0 <= left and right < len(s) and s[left] == s[right]:
if right - left + 1 > len(longest):
longest = s[left : right + 1]
left -= 1
right += 1
from std.testing import assert_true
def main() raises:
var s: String = "babad"
var expected: String = "bab"
result = longest_palindrome(s)
assert_true(result == expected, "Assertion failed")
s = "cbbd"
expected = "bb"
result = longest_palindrome(s)
assert_true(result == expected, "Assertion failed")
s = "racecar"
expected = "racecar"
result = longest_palindrome(s)
assert_true(result == expected, "Assertion failed")
Interleaving String
Determine if s3 is an interleaving of s1 and s2.
def is_interleave(s1: String, s2: String, s3: String) -> Bool:
if len(s1) + len(s2) != len(s3):
return False
dp = List[List[Bool]](
length=len(s1) + 1, fill=List[Bool](length=len(s2) + 1, fill=False)
)
dp[len(s1)][len(s2)] = True
for i in range(len(s1), -1, -1):
for j in range(len(s2), -1, -1):
if i < len(s1) and s1[i] == s3[i + j] and dp[i + 1][j]:
dp[i][j] = True
if j < len(s2) and s2[j] == s3[i + j] and dp[i][j + 1]:
dp[i][j] = True
return dp[0][0]
from std.testing import assert_true, assert_false
def main() raises:
var s1: String = "aabcc"
var s2: String = "dbbca"
var s3: String = "aadbbcbcac"
result = is_interleave(s1, s2, s3)
assert_true(result, "Assertion failed")
s1 = "aabcc"
s2 = "dbbca"
s3 = "aadbbbaccc"
result = is_interleave(s1, s2, s3)
assert_false(result, "Assertion failed")
s1 = ""
s2 = ""
s3 = ""
result = is_interleave(s1, s2, s3)
assert_true(result, "Assertion failed")
Last word length
Return the length of the last word in a given space-separated string
def last_word_length(s: String) -> Int:
if len(s) == 0:
return 0
i, length = len(s) - 1, 0
while i >= 0 and s[i] == " ":
i -= 1
while i >= 0 and s[i] != " ":
length += 1
i -= 1
return length
def main() raises:
from std.testing import assert_true
result = last_word_length("Hello World")
assert_true(result == 5, "Assertion failed")
result = last_word_length(" ")
assert_true(result == 0, "Assertion failed")
result = last_word_length(" fly me to the moon ")
assert_true(result == 4, "Assertion failed")
result = last_word_length("luffy is still joyboy")
assert_true(result == 6, "Assertion failed")
Find All Anagrams in a String.
Given two ASCII strings s and p, return a list of all start indices
of pβs anagrams in s. An anagram is a permutation of the characters
of p, so any contiguous substring of s whose character-frequency
dict matches that of p is a match.
Uses a fixed-length sliding-window β the window size equals len(p).
At each step the character-frequency dict of the window is compared
against the target dict (O(n) time, O(|Ξ£|) space).
Example:
s = "cbaebabacd", p = "abc" β [0, 6]
s = "abab", p = "ab" β [0, 1, 2]
from std.testing import assert_equal, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Sliding-window implementation
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def find_anagrams(s: String, p: String) -> List[Int]:
"""All start indices in `s` where a permutation of `p` occurs.
Builds a frequency dict for `p`, then slides a window of the same
length across `s`, maintaining a parallel frequency dict for the
current window. When the two dicts are equal the window is an
anagram match.
"""
var s_bytes = s.as_bytes()
var p_bytes = p.as_bytes()
var s_len = len(s_bytes)
var p_len = len(p_bytes)
if s_len < p_len or p_len == 0:
return List[Int]()
# Build frequency dicts for the first window.
var p_freq = Dict[UInt8, Int]()
var win_freq = Dict[UInt8, Int]()
for i in range(p_len):
p_freq[p_bytes[i]] = 1 + p_freq.get(p_bytes[i], 0)
win_freq[s_bytes[i]] = 1 + win_freq.get(s_bytes[i], 0)
var result = [0] if p_freq == win_freq else List[Int]()
var left: Int = 0
var right = p_len
# Slide the window one position at a time.
while right < s_len:
# Add the incoming character on the right.
win_freq[s_bytes[right]] = 1 + win_freq.get(s_bytes[right], 0)
# Remove the outgoing character on the left.
win_freq[s_bytes[left]] = win_freq.get(s_bytes[left], 1) - 1
# Clean up zero-count entries to keep dicts comparable.
if win_freq.get(s_bytes[left], 0) == 0:
_ = win_freq.pop(
s_bytes[left], -1
) # default -1 is ignored on success
right += 1
left += 1
if p_freq == win_freq:
result.append(left)
return result^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_basic_exact_match() raises:
assert_equal(find_anagrams("abc", "abc"), [0])
def test_reordered_match() raises:
assert_equal(find_anagrams("cba", "abc"), [0])
def test_extra_char_after() raises:
assert_equal(find_anagrams("cbaa", "abc"), [0])
def test_two_matches() raises:
assert_equal(find_anagrams("cbaacb", "abc"), [0, 3])
def test_cbaebabacd_example() raises:
assert_equal(find_anagrams("cbaebabacd", "abc"), [0, 6])
def test_abab_example() raises:
assert_equal(find_anagrams("abab", "ab"), [0, 1, 2])
def test_no_match() raises:
assert_equal(find_anagrams("abcdef", "xyz"), List[Int]())
def test_target_longer_than_source() raises:
assert_equal(find_anagrams("ab", "abc"), List[Int]())
def test_single_chars() raises:
assert_equal(find_anagrams("aaa", "a"), [0, 1, 2])
def test_single_char_no_match() raises:
assert_equal(find_anagrams("bbb", "a"), List[Int]())
def test_target_at_end() raises:
assert_equal(find_anagrams("xyzabc", "abc"), [3])
def test_target_at_start() raises:
assert_equal(find_anagrams("abcxyz", "abc"), [0])
def test_overlapping_windows() raises:
assert_equal(find_anagrams("aaaa", "aa"), [0, 1, 2])
def test_empty_source() raises:
assert_equal(find_anagrams("", "a"), List[Int]())
def test_empty_target() raises:
assert_equal(find_anagrams("abc", ""), List[Int]())
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Word Search
Check if a word can be formed in a grid by sequentially adjacent (non-repeating) horizontal or vertical letters.
from collections import Set
comptime SString = List[StaticString]
comptime SStrings = List[SString]
def present(board: SStrings, word: StaticString) raises -> Bool:
if len(board) == 0:
return False
rows, cols = len(board), len(board[0])
visited = Set[String]()
for row in range(rows):
for col in range(cols):
if trace(rows, cols, board, word, 0, row, col, visited):
return True
return False
def trace(
rows: UInt,
cols: UInt,
board: SStrings,
word: StaticString,
idx: UInt,
row: UInt,
col: UInt,
mut visited: Set[String],
) raises -> Bool:
if len(word) == idx:
return True
cell = String(row) + String(col)
if (
row < 0
or row >= rows
or col < 0
or col >= cols
or board[row][col] != word[idx]
or cell in visited
):
return False
visited.add(cell)
exists = (
trace(rows, cols, board, word, idx + 1, row + 1, col, visited)
or trace(rows, cols, board, word, idx + 1, row - 1, col, visited)
or trace(rows, cols, board, word, idx + 1, row, col + 1, visited)
or trace(rows, cols, board, word, idx + 1, row, col - 1, visited)
)
visited.remove(cell)
return exists
from std.testing import assert_true, assert_false
def main() raises:
board = SStrings(
SString("A", "B", "C", "E"),
SString("S", "F", "C", "S"),
SString("A", "D", "E", "E"),
)
word1 = "ABCB"
result = present(board, word1)
assert_false(result, "Assertion failed")
board = SStrings(
SString("A", "B", "C", "E"),
SString("S", "F", "C", "S"),
SString("A", "D", "E", "E"),
)
word2 = "SEE"
result = present(board, word2)
assert_true(result, "Assertion failed")
board = SStrings(
SString("A", "B", "C", "E"),
SString("S", "F", "C", "S"),
SString("A", "D", "E", "E"),
)
word3 = "ABCCED"
result = present(board, word3)
assert_true(result, "Assertion failed")
Valid Sudoku.
Determine if a 9 x 9 Sudoku board is valid. Only the filled cells need to be validated according to three rules:
- Each row must contain the digits 1β9 without repetition.
- Each column must contain the digits 1β9 without repetition.
- Each of the nine 3 x 3 sub-boxes must contain the digits 1β9 without repetition.
Note: a valid board is not necessarily solvable β only filled cells are checked.
Algorithm β O(81) = O(1) time and space:
Maintain three dictionaries mapping coordinates to a Set of seen
digits β one for rows (indexed by row number), one for columns
(indexed by column number), and one for 3x3 sub-boxes (keyed by
(row // 3, col // 3)). For each cell: if it contains a digit,
pop the corresponding sets, check for duplicates, add the digit,
and store the set back. If any duplicate is found the board is
invalid.
Example:
valid_sudoku([
["5","3",".",".","7",".",".",".","."],
["6",".",".","1","9","5",".",".","."],
[".","9","8",".",".",".",".","6","."],
["8",".",".",".","6",".",".",".","3"],
["4",".",".","8",".","3",".",".","1"],
["7",".",".",".","2",".",".",".","6"],
[".","6",".",".",".",".","2","8","."],
[".",".",".","4","1","9",".",".","5"],
[".",".",".",".","8",".",".","7","9"],
]) β true
(Swapping the top-left "5" for "8" makes the board invalid because
column 0 and the top-left sub-box would both contain duplicate 8s.)
from std.collections import Set
from std.testing import assert_false, assert_true, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Validator
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def valid_sudoku(board: List[List[String]]) -> Bool:
"""Check whether a 9x9 Sudoku board obeys the three validity rules.
One pass over the board: for each filled cell, the digit is checked
against the row-set, column-set, and sub-box-set for its position.
On first sight of a duplicate the function returns `False`.
"""
# Defensive shape check (a proper caller would always pass 9Γ9).
if len(board) != 9:
return False
for r in range(9):
if len(board[r]) != 9:
return False
# Three dictionaries tracking digits seen per row, column, and
# 3Γ3 sub-box. Each value is a Set that is "popped" (removed &
# returned), mutated, and then stored back.
var rows = Dict[Int, Set[String]]()
var cols = Dict[Int, Set[String]]()
var squares = Dict[Tuple[Int, Int], Set[String]]()
for r in range(9):
for c in range(9):
var value = board[r][c]
if value == ".":
continue
# Pop the three sets for the current position (or get a
# fresh empty set if this is the first cell in that group).
var row = rows.pop(r, Set[String]())
var col = cols.pop(c, Set[String]())
var square = squares.pop((r // 3, c // 3), Set[String]())
# Duplicate in any of the three groups β invalid.
if value in row or value in col or value in square:
return False
row.add(value)
col.add(value)
square.add(value)
# Store the mutated sets back into their dictionaries.
rows[r] = row^
cols[c] = col^
squares[(r // 3, c // 3)] = square^
return True
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# ββ valid boards ββββββββββββββββββββββββββββββββββββββββββββ
def test_valid_example() raises:
var board: List[List[String]] = [
["5", "3", ".", ".", "7", ".", ".", ".", "."],
["6", ".", ".", "1", "9", "5", ".", ".", "."],
[".", "9", "8", ".", ".", ".", ".", "6", "."],
["8", ".", ".", ".", "6", ".", ".", ".", "3"],
["4", ".", ".", "8", ".", "3", ".", ".", "1"],
["7", ".", ".", ".", "2", ".", ".", ".", "6"],
[".", "6", ".", ".", ".", ".", "2", "8", "."],
[".", ".", ".", "4", "1", "9", ".", ".", "5"],
[".", ".", ".", ".", "8", ".", ".", "7", "9"],
]
assert_true(valid_sudoku(board^))
def test_empty_board() raises:
var board: List[List[String]] = [
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
]
assert_true(valid_sudoku(board^))
def test_full_valid_board() raises:
var board: List[List[String]] = [
["5", "3", "4", "6", "7", "8", "9", "1", "2"],
["6", "7", "2", "1", "9", "5", "3", "4", "8"],
["1", "9", "8", "3", "4", "2", "5", "6", "7"],
["8", "5", "9", "7", "6", "1", "4", "2", "3"],
["4", "2", "6", "8", "5", "3", "7", "9", "1"],
["7", "1", "3", "9", "2", "4", "8", "5", "6"],
["9", "6", "1", "5", "3", "7", "2", "8", "4"],
["2", "8", "7", "4", "1", "9", "6", "3", "5"],
["3", "4", "5", "2", "8", "6", "1", "7", "9"],
]
assert_true(valid_sudoku(board^))
# ββ invalid boards ββββββββββββββββββββββββββββββββββββββββββ
def test_invalid_duplicate_in_row() raises:
var board: List[List[String]] = [
["8", "3", ".", ".", "7", ".", ".", ".", "."],
["6", ".", ".", "1", "9", "5", ".", ".", "."],
[".", "9", "8", ".", ".", ".", ".", "6", "."],
["8", ".", ".", ".", "6", ".", ".", ".", "3"],
["4", ".", ".", "8", ".", "3", ".", ".", "1"],
["7", ".", ".", ".", "2", ".", ".", ".", "6"],
[".", "6", ".", ".", ".", ".", "2", "8", "."],
[".", ".", ".", "4", "1", "9", ".", ".", "5"],
[".", ".", ".", ".", "8", ".", ".", "7", "9"],
]
# Two 8s in column 0 (rows 0 and 3) and in the top-left sub-box.
assert_false(valid_sudoku(board^))
def test_invalid_duplicate_in_column() raises:
var board: List[List[String]] = [
["5", "3", ".", ".", "7", ".", ".", ".", "."],
["6", ".", ".", "1", "9", "5", ".", ".", "."],
[".", "9", "8", ".", ".", ".", ".", "6", "."],
["8", ".", ".", ".", "6", ".", ".", ".", "3"],
["4", ".", ".", "8", ".", "3", ".", ".", "1"],
["7", ".", ".", ".", "2", ".", ".", ".", "6"],
[".", "6", ".", ".", ".", ".", "2", "8", "."],
[".", ".", ".", "4", "1", "9", ".", ".", "5"],
["5", ".", ".", ".", "8", ".", ".", "7", "9"],
]
# Two 5s in column 0 (rows 0 and 8).
assert_false(valid_sudoku(board^))
def test_invalid_duplicate_in_subbox() raises:
var board: List[List[String]] = [
["5", "3", ".", ".", "7", ".", ".", ".", "."],
["6", ".", ".", "1", "9", "5", ".", ".", "."],
[".", "9", "8", ".", ".", ".", ".", "6", "."],
["8", ".", ".", ".", "6", ".", ".", ".", "3"],
["4", ".", ".", "8", ".", "3", ".", ".", "1"],
["7", ".", ".", ".", "2", ".", ".", ".", "6"],
[".", "6", ".", ".", ".", ".", "2", "8", "."],
[".", ".", ".", "4", "1", "9", ".", ".", "5"],
[".", ".", ".", ".", "8", ".", ".", "7", "9"],
]
# Keep board from example 1 but insert duplicate 3 in top-left sub-box.
board[0][1] = String("3") # already 3 at (1,2), sub-box (0,0)
board[1][2] = String("3")
assert_false(valid_sudoku(board^))
def test_invalid_wrong_dimensions() raises:
var board = List[List[String]](capacity=3)
board.append(["1", "2", "3"])
board.append(["4", "5", "6"])
board.append(["7", "8", "9"])
assert_false(valid_sudoku(board^))
def test_invalid_all_same_digit() raises:
var board: List[List[String]] = [
["1", "1", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
[".", ".", ".", ".", ".", ".", ".", ".", "."],
]
assert_false(valid_sudoku(board^))
def test_invalid_digit_in_last_cell() raises:
"""A duplicate placed in the very last cell of a row/col/box."""
var board: List[List[String]] = [
["1", "2", "3", "4", "5", "6", "7", "8", "9"],
["4", "5", "6", "7", "8", "9", "1", "2", "3"],
["7", "8", "9", "1", "2", "3", "4", "5", "6"],
["2", "3", "4", "5", "6", "7", "8", "9", "1"],
["5", "6", "7", "8", "9", "1", "2", "3", "4"],
["8", "9", "1", "2", "3", "4", "5", "6", "7"],
["3", "4", "5", "6", "7", "8", "9", "1", "2"],
["6", "7", "8", "9", "1", "2", "3", "4", "5"],
["9", "1", "2", "3", "4", "5", "6", "7", "1"],
]
# Last cell duplicates 1 from the start of its row, column, and sub-box.
assert_false(valid_sudoku(board^))
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Encode and Decode Strings.
Design an algorithm to encode a list of ASCII strings into a single string and decode it back without loss. The challenge is to handle strings that may contain delimiters, digits, or any arbitrary content.
Length-prefix approach β format: <length>#<string> per element.
Encoding:
- For each string, compute its byte length.
- Append the length (as a decimal number), then
#, then the string.
Decoding:
- Scan for
#to locate the length prefix. - Parse the digits before
#as an integer β this is the byte length of the following string. - Extract exactly that many bytes after
#and add to the result. - Repeat until the entire encoded string is consumed.
This scheme is unambiguous because the length prefix tells the decoder exactly where each string ends, regardless of what characters it contains.
Example:
encode(["hello", "world"]) β "5#hello5#world"
decode("5#hello5#world") β ["hello", "world"]
from std.collections.string import Codepoint
from std.testing import assert_equal, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Encoder
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def encode(strs: List[String]) -> String:
"""Encode a list of strings into a single length-prefixed string.
For each input string, appends `byte_length` + `"#"` + the string
itself. The lengths are written as decimal ASCII digits so that
the decoder can always locate the next boundary via the `#` marker.
"""
var result = ""
for s in strs:
result += String(s.byte_length()) + "#" + s
return result^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Decoder
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def decode(s: String) raises -> List[String]:
"""Decode a length-prefixed string back into a list of strings.
1. Walk the byte span to find each `#` separator.
2. Read the digits before `#` as the byte-length of the string.
3. Slice exactly that many bytes after `#` and reconstruct the
original string.
"""
var result = List[String]()
var i = 0
comptime ascii_hash = UInt8(Int(Codepoint.ord("#")))
var bytes = s.as_bytes()
while i < len(bytes):
# Find the `#` separator β everything between i and j is the
# decimal length prefix.
var j = i
while j < len(bytes) and bytes[j] != ascii_hash:
j += 1
# Parse the length digits as an integer.
var length_digits = bytes[i:j]
var str_length = Int(String(from_utf8=length_digits))
# The string data starts right after `#` and occupies exactly
# `str_length` bytes.
var str_start = j + 1
result.append(String(from_utf8=bytes[str_start : str_start + str_length]))
# Advance past the string data to the start of the next entry.
i = str_start + str_length
return result^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_roundtrip_simple() raises:
var original: List[String] = ["hello", "world"]
assert_equal(decode(encode(original)), original)
def test_roundtrip_single_string() raises:
var original: List[String] = ["def"]
assert_equal(decode(encode(original)), original)
def test_roundtrip_empty_list() raises:
var original = List[String]()
assert_equal(decode(encode(original)), original)
def test_roundtrip_empty_string() raises:
var original: List[String] = [""]
assert_equal(decode(encode(original)), original)
def test_roundtrip_multiple_empty() raises:
var original: List[String] = ["", "", ""]
assert_equal(decode(encode(original)), original)
def test_roundtrip_with_delimiter() raises:
var original: List[String] = ["a#b", "c#d"]
assert_equal(decode(encode(original)), original)
def test_roundtrip_with_numbers() raises:
var original: List[String] = ["123", "456", "789"]
assert_equal(decode(encode(original)), original)
def test_roundtrip_mixed() raises:
var original: List[String] = ["def", "main()", "raises"]
assert_equal(decode(encode(original)), original)
def test_roundtrip_special_chars() raises:
var original: List[String] = ["!@#$%", " spaces ", "\t\n"]
assert_equal(decode(encode(original)), original)
def test_roundtrip_long_string() raises:
var original: List[String] = [
"a" * 100,
"b" * 200,
]
assert_equal(decode(encode(original)), original)
def test_roundtrip_all_ascii_printable() raises:
var all_chars = ""
for i in range(32, 127): # printable ASCII range
all_chars += String(chr(i))
var original: List[String] = [all_chars]
assert_equal(decode(encode(original)), original)
def test_roundtrip_preserves_order() raises:
var original: List[String] = [
"first", "second", "third", "fourth", "fifth",
]
var decoded = decode(encode(original))
assert_equal(len(decoded), 5)
for i in range(5):
assert_equal(decoded[i], original[i])
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Reverse String.
Reverse a string by iterating its Unicode codepoints in reverse order.
Algorithm β O(n) time and space.
- Collect all codepoints into a list.
- Iterate from the last index down to 0, appending each codepoint
to a result
String.
Example:
reverse("abc") β "cba"
reverse("") β ""
from std.testing import assert_equal, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Implementation
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def reverse(s: String) -> String:
"""Return `s` with its Unicode codepoints reversed.
Works on full Unicode codepoints (not raw bytes), so multi-byte
characters are preserved correctly.
"""
var codepoints = List[Codepoint](capacity=s.count_codepoints())
for code in s.codepoints():
codepoints.append(code)
if len(codepoints) == 0 or len(codepoints) == 1:
return s
var result = String()
for i in range(len(codepoints) - 1, -1, -1):
result.append(codepoints[i])
return result
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_empty_string() raises:
assert_equal(reverse(""), "")
def test_single_char() raises:
assert_equal(reverse("a"), "a")
def test_two_chars() raises:
assert_equal(reverse("ab"), "ba")
def test_three_chars() raises:
assert_equal(reverse("abc"), "cba")
def test_palindrome() raises:
assert_equal(reverse("racecar"), "racecar")
def test_with_spaces() raises:
assert_equal(reverse("hello world"), "dlrow olleh")
def test_unicode_multi_byte() raises:
assert_equal(reverse("cafΓ©"), "Γ©fac")
def test_unicode_emoji() raises:
assert_equal(reverse("aπb"), "bπa")
def test_numbers_and_symbols() raises:
assert_equal(reverse("123!@#"), "#@!321")
def test_reverse_twice_returns_original() raises:
var original = "Hello, δΈη! π"
assert_equal(reverse(reverse(original)), original)
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Reverse Vowels of a String.
Given a string s, reverse only the vowels in the string and return it.
Vowels are a, e, i, o, u in both lower and upper case.
Algorithm β O(n) time, O(n) space:
Use two pointers (left, right) starting at both ends of the
codepoint list. Advance each pointer inward until it lands on a
vowel, then swap the two vowels and continue. When the pointers
cross, all vowels have been reversed while consonants stay in
place.
Example:
s = "IceCreAm" β "AceCreIm"
s = "leetcode" β "leotcede"
from std.collections import Set
from std.collections.string import Codepoint
from std.testing import assert_equal, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Implementation
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def reverse_vowels(s: String) -> String:
"""Reverse only the vowels in `s`, leaving consonants in place."""
var codepoints = [code for code in s.codepoints()]
var n = len(codepoints)
if n == 0 or n == 1:
return s
var vowels = Set([Int(code) for code in "aeiouAEIOU".codepoints()])
var left = 0
var right = n - 1
while left < right:
# Advance left until it lands on a vowel.
while left < right and Int(codepoints[left]) not in vowels:
left += 1
# Advance right until it lands on a vowel.
while left < right and Int(codepoints[right]) not in vowels:
right -= 1
if left < right:
codepoints.swap_elements(left, right)
left += 1
right -= 1
var result = String()
for code in codepoints:
result.append(code)
return result
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_example_1() raises:
assert_equal(reverse_vowels("IceCreAm"), "AceCreIm")
def test_example_2() raises:
assert_equal(reverse_vowels("leetcode"), "leotcede")
def test_empty_string() raises:
assert_equal(reverse_vowels(""), "")
def test_single_char_consonant() raises:
assert_equal(reverse_vowels("b"), "b")
def test_single_char_vowel() raises:
assert_equal(reverse_vowels("a"), "a")
def test_no_vowels() raises:
assert_equal(reverse_vowels("bcdfg"), "bcdfg")
def test_all_vowels() raises:
assert_equal(reverse_vowels("aeiou"), "uoiea")
def test_vowels_in_mixed_case() raises:
assert_equal(reverse_vowels("AEIOU"), "UOIEA")
def test_only_middle_vowel() raises:
assert_equal(reverse_vowels("abcd"), "abcd")
def test_multiple_vowels_odd_length() raises:
assert_equal(reverse_vowels("hello"), "holle")
def test_vowels_at_ends() raises:
assert_equal(reverse_vowels("amazing"), "imazang")
def test_palindrome_with_vowels() raises:
assert_equal(reverse_vowels("racecar"), "racecar")
def test_spaces_and_punctuation() raises:
assert_equal(reverse_vowels("a!e@i#o$u%"), "u!o@i#e$a%")
def test_uppercase_and_lowercase() raises:
# Vowels are reversed as-is; case stays with each character.
assert_equal(reverse_vowels("AaEeIiOoUu"), "uUoOiIeEaA")
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Minimum Window Substring.
Given a string s and a target string t, find the shortest contiguous
substring of s that contains every character from t (including
multiplicity). If no such window exists, return None.
This file provides two implementations:
β’ min_window β O(n) sliding-window algorithm (primary solution).
β’ min_window_bruteforce β O(nΒ³) reference that checks every possible
window (kept for correctness verification).
Example:
s = "ADOBECODEBANC", t = "ABC"
β "BANC"
from std.testing import assert_equal, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# O(n) sliding-window solution
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def min_window(s: String, t: String) raises -> Optional[String]:
"""Shortest substring of `s` containing all chars in `t` (sliding window)."""
var source_bytes = s.as_bytes()
var target_bytes = t.as_bytes()
var source_length = len(source_bytes)
var target_length = len(target_bytes)
if source_length == 0 or target_length == 0 or source_length < target_length:
return None
# Build frequency dict for target characters.
var target_freq = Dict[UInt8, Int]()
for ch in target_bytes:
target_freq[ch] = target_freq.get(ch, 0) + 1
var required = len(target_freq)
# ββ sliding-window loop ββββββββββββββββββββββββββββββββββ
var window_freq = Dict[UInt8, Int]()
var formed = 0
var left: Int = 0
var shortest_window_length: Int = source_length + 1
var shortest_window_start: Int = 0
for right in range(source_length):
# Expand window by one character on the right.
var ch = source_bytes[right]
window_freq[ch] = window_freq.get(ch, 0) + 1
if ch in target_freq and window_freq[ch] == target_freq[ch]:
formed += 1
# Contract from the left while the window is still valid.
while formed == required:
var current_window_length = right - left + 1
if current_window_length < shortest_window_length:
shortest_window_length = current_window_length
shortest_window_start = left
# Remove leftmost character from the window.
var left_char = source_bytes[left]
window_freq[left_char] = window_freq[left_char] - 1
if left_char in target_freq and window_freq[left_char] < target_freq[left_char]:
formed -= 1
left += 1
if shortest_window_length <= source_length:
var sub_bytes = source_bytes[
shortest_window_start : shortest_window_start + shortest_window_length
]
return String(from_utf8=sub_bytes)
return None
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# O(nΒ³) brute-force reference (for correctness verification)
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def min_window_bruteforce(s: String, t: String) raises -> Optional[String]:
"""Shortest substring of `s` containing all chars in `t` (brute force).
Exhaustively checks every possible substring. Used to verify the
O(n) sliding-window implementation.
"""
var source_bytes = s.as_bytes()
var target_bytes = t.as_bytes()
var source_length = len(source_bytes)
var target_length = len(target_bytes)
if source_length == 0 or target_length == 0 or source_length < target_length:
return None
var target_frequencies = Dict[UInt8, Int]()
for ch in target_bytes:
target_frequencies[ch] = target_frequencies.get(ch, 0) + 1
var unique_target_char_count = len(target_frequencies)
var shortest_window_length: Int = source_length + 1
var shortest_window_start: Int = 0
var window_frequencies = Dict[UInt8, Int]()
for window_start in range(source_length - target_length + 1):
if source_bytes[window_start] not in target_frequencies:
continue
for window_end in range(window_start + target_length - 1, source_length):
if window_end == window_start + target_length - 1:
window_frequencies.clear()
for ch in source_bytes[window_start : window_end + 1]:
window_frequencies[ch] = window_frequencies.get(ch, 0) + 1
else:
var ch = source_bytes[window_end]
window_frequencies[ch] = window_frequencies.get(ch, 0) + 1
var satisfied_char_types = 0
for item in target_frequencies.items():
var letter = item.key
var required_count = item.value
if window_frequencies.get(letter, 0) >= required_count:
satisfied_char_types += 1
if satisfied_char_types == unique_target_char_count:
var current_window_length = window_end - window_start + 1
if current_window_length < shortest_window_length:
shortest_window_length = current_window_length
shortest_window_start = window_start
break
if shortest_window_length <= source_length:
var shortest_substring_bytes = source_bytes[
shortest_window_start : shortest_window_start + shortest_window_length
]
return String(from_utf8=shortest_substring_bytes)
return None
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_example_1() raises:
assert_equal(min_window("ADOBECODEBANC", "ABC"), "BANC")
def test_example_2() raises:
assert_equal(min_window("a", "a"), "a")
def test_example_3() raises:
assert_equal(min_window("a", "aa"), None)
def test_contains_ain() raises:
assert_equal(min_window("contains", "ain"), "ain")
def test_shortest_ntai() raises:
assert_equal(min_window("ntai", "ain"), "ntai")
def test_negative_nums_matched() raises:
assert_equal(min_window("xaybz", "ab"), "ayb")
def test_equal_length() raises:
assert_equal(min_window("abc", "abc"), "abc")
def test_no_match() raises:
assert_equal(min_window("a", "b"), None)
def test_duplicates_in_target() raises:
assert_equal(min_window("aa", "aa"), "aa")
assert_equal(min_window("aba", "aa"), "aba")
def test_window_at_end() raises:
assert_equal(min_window("bac", "ac"), "ac")
def test_full_string_is_only_window() raises:
assert_equal(min_window("abcdef", "az"), None)
assert_equal(min_window("abcdef", "fed"), "def")
def test_bf_matches_sliding_window() raises:
"""Verify that both implementations agree on a diverse set of inputs."""
var cases = List[String]()
var targets = List[String]()
cases.append("ADOBECODEBANC")
targets.append("ABC")
cases.append("figehaeci")
targets.append("aei")
cases.append("contains")
targets.append("ain")
cases.append("xaybz")
targets.append("ab")
cases.append("abc")
targets.append("abc")
cases.append("aa")
targets.append("aa")
cases.append("aba")
targets.append("aa")
cases.append("bac")
targets.append("ac")
cases.append("abcdef")
targets.append("fed")
for i in range(len(cases)):
var sliding = min_window(cases[i], targets[i])
var brute = min_window_bruteforce(cases[i], targets[i])
assert_equal(sliding, brute, "Mismatch between implementations")
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Buy And Sell Stock
# Function to calculate the maximum profit from a list of stock prices
# where you are allowed to make only one buy and one sell transaction.
# You must buy before you sell.
def max_profit(prices: List[UInt]) -> UInt:
var max_profit: UInt = 0 # Stores the maximum profit found so far
buy_day = 0 # Pointer to track the day to buy the stock
sell_day = 1 # Pointer to track the day to sell the stock
# Loop until sell_day reaches the end of the price list
while sell_day < len(prices):
if prices[buy_day] < prices[sell_day]:
# If selling is profitable, calculate profit and update max
max_profit = max(max_profit, prices[sell_day] - prices[buy_day])
else:
# If current sell_day is cheaper than buy_day, shift buy_day
buy_day = sell_day
sell_day += 1 # Move to the next day
return max_profit # Return the highest profit found
def main():
# First test: Best profit is buying at 1 and selling at 6 => profit = 5
prices = [7, 1, 5, 3, 6, 4]
debug_assert(max_profit(prices) == 5, "Assertion failed")
# Second test: No profitable day to sell => profit = 0
prices = [7, 6, 4, 3, 1]
debug_assert(max_profit(prices) == 0, "Assertion failed")
Buy And Sell Stock 2
# On each day, you may decide to buy and/or sell the stock. You can only hold
# at most one share of the stock at any time. However, you can buy it then
# immediately sell it on the same day.
# Find and return the maximum profit you can achieve.
# Note the problem is confusing - given per day prices, you make no profit buying
# and selling on the same day!
def total_profit(prices: List[UInt]) -> UInt:
var total_profit: UInt = 0 # Stores the maximum profit found so far
# Loop until sell day reaches the end of the price list
for day in range(1, len(prices)):
if prices[day - 1] < prices[day]:
# If selling is profitable, sell it & add up profit
total_profit += prices[day] - prices[day - 1]
return total_profit # Return the highest profit found
def main():
# First test: Best profit is buying at 1 and selling at 6 => profit = 5
prices = [7, 1, 5, 3, 6, 4]
debug_assert(total_profit(prices) == 7, "Assertion failed")
# Second test: No profitable day to sell => profit = 0
prices = [7, 6, 4, 3, 1]
debug_assert(total_profit(prices) == 0, "Assertion failed")
# 3rd test: Best profit is buying at every day and selling next day
prices = [1, 2, 3, 4, 5]
debug_assert(total_profit(prices) == 4, "Assertion failed")
Min In Sorted Rotated Arr
# Given a sorted, rotated array `nums`, find and return the minimum element.
# The solution uses binary search to achieve O(log n) time complexity.
def find_min(nums: List[Int]) -> Int:
# Handle edge case: empty list
if len(nums) == 0:
return Int.MIN # Return the minimum representable integer
# Handle edge case: single-element list
elif len(nums) == 1:
return nums[0]
else:
# Initialize binary search pointers
left, right = 0, len(nums) - 1
curr_min = nums[0] # Assume first element is the minimum initially
while left <= right:
# If the current window is sorted, the smallest element is at the left
if nums[left] <= nums[right]:
return min(curr_min, nums[left])
# Compute the mid index
mid = (left + right) // 2
# Update the minimum seen so far
cur_min = min(curr_min, nums[mid])
# Determine which side is unsorted (contains the pivot)
if nums[mid] >= nums[left]:
# Left half is sorted, so min must be in the right half
left = mid + 1
else:
# Right half is sorted, so min must be in the left half (including mid)
right = mid - 1
# Fallback return β should never be reached in a rotated sorted array
return curr_min # Keeps compiler happy
def main():
nums = List[Int]()
minimum = find_min(nums)
debug_assert(minimum == Int.MIN, "Assertion failed")
nums = [4, 5, 6, 7, 0, 1, 2]
minimum = find_min(nums)
debug_assert(minimum == 0, "Assertion failed")
nums = [4]
minimum = find_min(nums)
debug_assert(minimum == 4, "Assertion failed")
nums = [4, 5]
minimum = find_min(nums)
debug_assert(minimum == 4, "Assertion failed")
nums = [5, 4]
minimum = find_min(nums)
debug_assert(minimum == 4, "Assertion failed")
nums = [3, 4, 5, 1, 2]
minimum = find_min(nums)
debug_assert(minimum == 1, "Assertion failed")
nums = [11, 13, 15, 17]
minimum = find_min(nums)
debug_assert(minimum == 11, "Assertion failed")
nums = [11, 13, 15, 17, 1, 1, 2, 2]
minimum = find_min(nums)
debug_assert(minimum == 1, "Assertion failed")
Search Sorted Rotated Arr
# Search in Rotated Sorted Array
# Function to search for a target in a rotated sorted array
comptime ItemType = ComparableCollectionElement
def find[ItmType: ItemType](read items: List[ItmType], target: ItmType) -> Int:
if len(items) == 0:
return -1
# Initialize pointers for binary search
left, right = 0, len(items) - 1
# Perform binary search
while left <= right:
mid = (left + right) // 2 # Calculate middle index
# If the middle element is the target, return the index
if items[mid] == target:
return mid
# Determine which half is sorted
if items[mid] >= items[left]:
# Left half is sorted
# Check if target lies outside the sorted left half
if target < items[left] or target > items[mid]:
# Target is in the right half
left = mid + 1
else:
# Target is in the left half
right = mid - 1
else:
# Right half is sorted
# Check if target lies outside the sorted right half
if target > items[right] or target < items[mid]:
# Target is in the left half
right = mid - 1
else:
# Target is in the right half
left = mid + 1
# Target not found
return -1
def main():
# Example 1: Target exists in the array
items = [4, 5, 6, 7, 0, 1, 2]
target = 0
# Expected output: 4 (index of 0)
# debug_assert(find(items, target) == 4, "Assertion failed")
# Example 2: Target does not exist
items = [4, 5, 6, 7, 0, 1, 2]
target = 3
# Expected output: -1
# debug_assert(find(items, target) == -1, "Assertion failed")
# Example 3: Single-element array, target not present
items = [1]
target = 0
# Expected output: -1
debug_assert(find(items, target) == -1, "Assertion failed")
Search in a rotated sorted array that may contain duplicates
Generic function to find the index of a target element in a rotated sorted list. Works for any type that implements ComparableCollectionElement (e.g., Int, Float, etc.).
# Define an alias for types that support comparison operations.
comptime ItemType = ComparableCollectionElement
def find[ItmType: ItemType](read items: List[ItmType], target: ItmType) -> Int:
if len(items) == 0:
return -1
left, right = 0, len(items) - 1
# Perform modified binary search to handle rotation and duplicates
while left <= right:
mid = left + (right - left) // 2
# Target found at midpoint
if items[mid] == target:
return mid
# Case 1: Target is less than midpoint value
if target < items[mid]:
# If target is greater than the left bound, it must be in the left subarray
if target > items[left]:
right = mid - 1
# If target is less than the left bound, it must be in the right subarray
elif target < items[left]:
left = mid + 1
# If target equals the left bound, it's a match
else:
return left
# Case 2: Target is greater than midpoint value
else:
# If target is less than the right bound, it lies in the right subarray
if target < items[right]:
left = mid + 1
# If target is greater than the right bound, it must lie to the left
elif target > items[right]:
right = mid - 1
# If target equals the right bound, it's a match
else:
return right
# Target not found
return -1
from std.testing import assert_equal
def main() raises:
items = [7, 7, 8, 9, 10, 10, 12, 1, 2, 3, 3, 4, 4, 5, 5, 6]
targets = [12, 7, 3, 10, 1, 6, 4]
expected_indices = [6, 0, 9, 5, 7, 15, 11]
results = List[Int](capacity=len(targets))
for i in range(len(targets)):
index = find(items, targets[i])
results.append(index)
assert_equal(expected_indices, results, "Assertion failed!")
items = [6]
target = 6
index = find(items, target)
assert_equal(index, 0, "Assertion failed")
items = [7, 7, 7]
target = 7
index = find(items, target)
assert_equal(index, 1, "Assertion failed")
items = [7, 8, 4, 5]
target = 3
index = find(items, target)
assert_equal(index, -1, "Assertion failed")
Find First/Last
Find first and last index of a target value in a sorted array
def find_first_last(arr: List[Int], target: Int) -> (Int, Int):
result = (-1, -1)
if len(arr) == 0:
return result
left, right = 0, len(arr) - 1
while left <= right:
mid = (left + right) // 2
if arr[mid] == target:
result[1] = mid
left = mid + 1
elif arr[mid] > target:
right = mid - 1
else:
left = mid + 1
left, right = (
0,
result[1],
) # result[1] -1 would keep left index at -1 for single occurence of target
while left <= right:
mid = (left + right) // 2
if arr[mid] == target:
result[0] = mid
right = mid - 1
elif arr[mid] > target:
right = mid - 1
else:
left = mid + 1
return result
from std.testing import assert_true
def main() raises:
arr = [5, 7, 7, 8, 8, 10]
target = 8
result = find_first_last(arr, target)
assert_true(result[0] == 3 and result[1] == 4, "Assertion failed")
target = 6
result = find_first_last(arr, target)
assert_true(result[0] == -1 and result[1] == -1, "Assertion failed")
arr = [5, 7, 7, 8, 10]
target = 8
result = find_first_last(arr, target)
assert_true(result[0] == 3 and result[1] == 3, "Assertion failed")
Merge Itervals
Merge overlapping intervals
@parameter
def compare_fn(interval1: (Int, Int), interval2: (Int, Int)) -> Bool:
return interval1[0] < interval2[0]
def merge_intervals(mut intervals: List[(Int, Int)]) -> List[(Int, Int)]:
if len(intervals) == 0:
return List[(Int, Int)]()
sort[compare_fn](intervals)
result = List[(Int, Int)]()
result.append(intervals[0])
for curr_interval in intervals[1:]:
start, end = curr_interval[]
last_interval = result[-1]
last_start, last_end = last_interval
if start <= last_end:
result[len(result) - 1] = (last_start, max(last_end, end))
else:
result.append(curr_interval[])
return result
from std.testing import assert_true
def main() raises:
intervals = List[(Int, Int)]((1, 3), (2, 6), (8, 10), (15, 18))
expected = List[(Int, Int)]((1, 6), (8, 10), (15, 18))
result = merge_intervals(intervals)
i = 0
for each in result:
assert_true(
each[][0] == expected[i][0] and each[][1] == expected[i][1],
"Assertion failed",
)
i += 1
intervals = List[(Int, Int)]((1, 4), (4, 5))
expected = List[(Int, Int)]((1, 5))
result = merge_intervals(intervals)
i = 0
for each in result:
assert_true(
each[][0] == expected[i][0] and each[][1] == expected[i][1],
"Assertion failed",
)
i += 1
Shapes
trait Shape(ComparableCollectionElement):
def area(self) -> UInt:
...
@value
struct Rectangle(Shape):
var length: UInt
var width: UInt
def area(self) -> UInt:
return self.length * self.width
def __lt__(self, other: Self) -> Bool:
return self.area() < other.area()
def __le__(self, other: Self) -> Bool:
return self.area() <= other.area()
def __eq__(self, other: Self) -> Bool:
return self.area() == other.area()
def __ne__(self, other: Self) -> Bool:
return self.area() != other.area()
def __gt__(self, other: Self) -> Bool:
return self.area() > other.area()
def __ge__(self, other: Self) -> Bool:
return self.area() >= other.area()
Generic singnly link list
A singly linked list with parametric polymorphism.
Current supports adding multiple elements at one go via the append method
from memory import Pointer, UnsafePointer
comptime ElementType = CollectionElement
@value
struct Node[
T: ElementType,
]:
comptime NextNode = UnsafePointer[Self]
var value: T
var next: Self.NextNode
def __init__(
out self,
owned value: T,
):
self.value = value
self.next = Self.NextNode()
def __init__(
out self,
owned value: T,
next: Optional[Self.NextNode],
):
self.value = value^
self.next = next.value() if next else Self.NextNode()
def __bool__(self) -> Bool:
return True
def __str__[
ElementType: WritableCollectionElement
](self: Node[ElementType]) -> String:
return String.write(self.value)
struct LinkedList[T: ElementType](Sized):
var head: Optional[Node[T]]
var len: UInt
def __init__(out self):
self.head = None
self.len = 0
def __len__(self) -> Int:
return self.len
def __init__(out self, *elems: T):
self = Self()
self.append(elems)
def append(mut self, *elems: T):
self.append(elems)
def append(mut self, elems: VariadicListMem[T]):
if len(elems) == 0:
return
next = 0
var current: UnsafePointer[Node[T]]
if self.head is None:
self.head = Optional(Node(elems[0]))
current = UnsafePointer(to=self.head.value())
next = 1
self.len += 1
else:
curr = UnsafePointer(to=self.head.value())
while curr and curr[].next:
curr = curr[].next
current = curr
for i in range(next, len(elems)):
node = Node(elems[i])
current[].next = UnsafePointer[Node[T]].alloc(1)
current[].next.init_pointee_move(node)
current = current[].next
self.len += 1
def __str__[
ElementType: WritableCollectionElement
](self: LinkedList[ElementType]) -> String:
if self.len == 0:
return String("[]")
else:
s = String("[")
current = self.head.value()
s.write(current.value)
for i in range(1, self.len):
next = current.next[]
if i <= self.len - 1:
s.write(", ")
s.write(next.value)
current = next
s.write("]")
return s
def __iter__(self) -> _LinkedListIter[T, __origin_of(self)]:
return _LinkedListIter(Pointer(to=self))
@value
struct _LinkedListIter[
mut: Bool, //,
ElementType: CollectionElement,
origin: Origin[mut],
]:
var src: Pointer[LinkedList[ElementType], origin]
var curr: UnsafePointer[Node[ElementType]]
var moved: Int
def __init__(out self, src: Pointer[LinkedList[ElementType], origin]):
self.src = src
self.curr = UnsafePointer(to=self.src[].head.value())
self.moved = 0
def __itr__(self) -> Self:
return self
def __next__(mut self) -> Pointer[ElementType, origin]:
out = Pointer[ElementType, origin](to=self.curr[].value)
self.moved += 1
self.curr = self.curr[].next
return out
def __has_next__(self) -> Bool:
return self.curr.__bool__()
def __len__(self) -> Int:
return self.src[].len - self.moved
def main():
linkedlist = LinkedList[Int]()
print(linkedlist.__str__())
linkedlist = LinkedList(1)
print(linkedlist.__str__())
linkedlist = LinkedList(1, 2, 3)
print(linkedlist.__str__())
linkedlist.append(4, 5, 6)
print(linkedlist.len)
print(linkedlist.__str__())
for e in linkedlist:
print(e[].__str__())
Longest Common Subsequence Recursive
Return the length of the longest common subsequence between two strings, or 0 if none exists
def longest_subseq(mut text1: String, mut text2: String) raises -> Int:
if len(text1) == 0 or len(text2) == 0:
return 0
if text1[len(text1) - 1] == text2[len(text2) - 1]:
text1 = text1[0:-1]
text2 = text2[0:-1]
return 1 + longest_subseq(text1, text2)
else:
text_1 = text1[0:-1]
text_2 = text2[0:-1]
count1 = longest_subseq(text1, text_2)
count2 = longest_subseq(text2, text_1)
return max(count1, count2)
from std.testing import assert_equal
def main() raises:
var text1: String = "abcde"
var text2: String = "ace"
result = longest_subseq(text1, text2)
assert_equal(result, 3, "Assertion failed")
text1 = "abc"
text2 = "abc"
result = longest_subseq(text1, text2)
assert_equal(result, 3, "Assertion failed")
text1 = "abc"
text2 = "xyz"
result = longest_subseq(text1, text2)
assert_equal(result, 0, "Assertion failed")
Longest Common Subsequence Dynamic
Return the length of the longest common subsequence between two strings, or 0 if none exists
def longest_subseq(mut text1: String, mut text2: String) raises -> Int:
if len(text1) == 0 or len(text2) == 0:
return 0
dp = List[List[Int]](
length=len(text1) + 1, fill=List[Int](length=len(text2) + 1, fill=0)
)
for i in range(len(text1) - 1, -1, -1):
for j in range(len(text2) - 1, -1, -1):
if text1[i] == text2[j]:
dp[i][j] = 1 + dp[i + 1][j + 1]
else:
dp[i][j] = max(dp[i][j + 1], dp[i + 1][j])
return dp[0][0]
from std.testing import assert_equal
def main() raises:
var text1: String = "abcde"
var text2: String = "ace"
result = longest_subseq(text1, text2)
assert_equal(result, 3, "Assertion failed")
text1 = "abc"
text2 = "abc"
result = longest_subseq(text1, text2)
assert_equal(result, 3, "Assertion failed")
text1 = "abc"
text2 = "xyz"
result = longest_subseq(text1, text2)
assert_equal(result, 0, "Assertion failed")
Combination sum
Find all unique combinations of numbers from candidates (reusable unlimited times) that sum to target.
def combination_sum(candidates: List[Int], target: Int) -> List[List[Int]]:
combinations = List[List[Int]]()
if len(candidates) == 0:
return combinations
var curr_combination = []
find_combinations(candidates, 0, curr_combination, 0, target, combinations)
return combinations
def find_combinations(
candidates: List[Int],
curr_index: Int,
mut curr_combination: List[Int],
total: Int,
target: Int,
mut combinations: List[List[Int]],
):
if total == target:
copy = curr_combination.copy()
sort(copy) # For validation
combinations.append(copy)
return
if curr_index >= len(candidates) or total > target:
return
curr_combination.append(candidates[curr_index])
find_combinations(
candidates,
curr_index,
curr_combination,
total + candidates[curr_index],
target,
combinations,
)
_ = curr_combination.pop()
find_combinations(
candidates,
curr_index + 1,
curr_combination,
total,
target,
combinations,
)
from std.testing import assert_true
def main() raises:
candidates = [2, 3, 6, 7]
target = 7
var result: List[List[Int]] = combination_sum(candidates, target)
expected = [2, 2, 3]
count = 0
for each in result:
if each[] == expected:
count += 1
assert_true(count == 1, "assertion failed")
expected = [7]
count = 0
for each in result:
if each[] == expected:
count += 1
assert_true(count == 1, "assertion failed")
candidates = [2, 3, 5]
target = 8
result = combination_sum(candidates, target)
expected = [2, 2, 2, 2]
count = 0
for each in result:
if each[] == expected:
count += 1
assert_true(count == 1, "assertion failed")
expected = [2, 3, 3]
count = 0
for each in result:
if each[] == expected:
count += 1
assert_true(count == 1, "assertion failed")
expected = [3, 5]
count = 0
for each in result:
if each[] == expected:
count += 1
assert_true(count == 1, "assertion failed")
candidates = [2]
target = 1
result = combination_sum(candidates, target)
assert_true(len(result) == 0, "assertion failed")
Permutations
Given an array of distinct integers nums, return all possible
permutations. The answer may be returned in any order.
Algorithm β O(n Β· n!) time, O(n Β· n!) space (output size).
Uses recursive backtracking:
1. For each element in nums, pop it, recursively compute all
permutations of the remaining nβ1 elements, then append the
popped element to each sub-permutation.
2. Restore nums by appending the element back before the next
iteration (backtracking step).
3. Base case: a single-element list has exactly one permutation.
The total number of permutations is n!. The result list is pre-allocated to this capacity to avoid repeated resizing.
Example:
permute([1, 2, 3]) β [[1, 2, 3], [1, 3, 2], [2, 1, 3],
[2, 3, 1], [3, 1, 2], [3, 2, 1]]
permute([0, 1]) β [[0, 1], [1, 0]]
permute([1]) β [[1]]
from std.testing import assert_equal, assert_true, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Implementation
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def permute(mut nums: List[Int]) -> List[List[Int]]:
"""Return all permutations of `nums` via recursive backtracking.
The input list is mutated during recursion and restored before
returning (the caller sees the original order).
"""
if len(nums) == 1:
return [nums.copy()]
# Pre-compute capacity: n! is the exact number of permutations.
var capacity = 1
for i in range(2, len(nums) + 1):
capacity *= i
var result = List[List[Int]](capacity=capacity)
for _ in range(len(nums)):
# Take one element from the front.
var n = nums.pop(0)
# Recursively permute the remaining elements.
var perms = permute(nums)
# Append the removed element to every sub-permutation.
for ref perm in perms:
perm.append(n)
# Collect all results and restore `nums` (backtrack).
result.extend(perms^)
nums.append(n)
return result^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Helpers
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def _is_permutation_of(perm: List[Int], original: List[Int]) -> Bool:
"""Check that `perm` contains exactly the elements of `original`."""
if len(perm) != len(original):
return False
for x in original:
if x not in perm:
return False
return True
def _all_permutations_are_unique(perms: List[List[Int]]) -> Bool:
"""Check that no two permutations in the list are identical."""
for i in range(len(perms)):
for j in range(i + 1, len(perms)):
if perms[i] == perms[j]:
return False
return True
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_example_1() raises:
var nums: List[Int] = [1, 2, 3]
var result = permute(nums)
assert_equal(len(result), 6)
for p in result:
assert_true(_is_permutation_of(p, [1, 2, 3]))
assert_true(_all_permutations_are_unique(result))
def test_example_2() raises:
var nums: List[Int] = [0, 1]
var result = permute(nums)
assert_equal(len(result), 2)
for p in result:
assert_true(_is_permutation_of(p, [0, 1]))
assert_true(_all_permutations_are_unique(result))
def test_single_element() raises:
var nums: List[Int] = [42]
var result = permute(nums)
assert_equal(len(result), 1)
assert_equal(result[0], [42])
def test_four_elements() raises:
var nums: List[Int] = [1, 2, 3, 4]
var result = permute(nums)
assert_equal(len(result), 24)
for p in result:
assert_true(_is_permutation_of(p, [1, 2, 3, 4]))
assert_true(_all_permutations_are_unique(result))
def test_two_elements_reversed() raises:
var nums: List[Int] = [5, 10]
var result = permute(nums)
assert_equal(len(result), 2)
assert_true(_all_permutations_are_unique(result))
def test_input_unchanged_after_call() raises:
var nums: List[Int] = [1, 2, 3]
var copy = nums.copy()
_ = permute(nums)
# The function should restore `nums` to its original state.
assert_equal(nums, copy)
def test_negative_numbers() raises:
var nums: List[Int] = [-1, 0, 1]
var result = permute(nums)
assert_equal(len(result), 6)
for p in result:
assert_true(_is_permutation_of(p, [-1, 0, 1]))
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Permutations of a String
Given an ASCII string word, return all possible permutations of its
characters. The answer may be returned in any order.
Two implementations are provided:
-
permute_recursiveβ O(n Β· n!) recursive backtracking. For each character position, permute the remaining substring recursively, then append the selected character to each sub-permutation. -
permute_iterativeβ O(n Β· n!) incremental insertion. Start with the last character, then repeatedly insert the next character at every possible position in every existing permutation.
Example:
permute_recursive("abc") β [bca, cba, acb, cab, abc, bac] (order varies)
permute_iterative("abc") β [abc, bac, bca, acb, cab, cba] (order varies)
from std.collections import Set
from std.collections.string import Codepoint
from std.testing import assert_equal, assert_true, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Approach 1 β recursive backtracking
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def permute_recursive(word: String) -> List[String]:
"""Return all permutations of `word` via recursion.
For each character at index `i`, build the string without that
character (`prefix + suffix`), recursively find all permutations
of that smaller string, then append character `i` to each result.
"""
var codepoints = [codepoint for codepoint in word.codepoints()]
var length = len(codepoints)
if length == 0 or length == 1:
return [word]
var capacity = 1
for i in range(2, length + 1):
capacity *= i
var perms = List[String](capacity=capacity)
for i in range(length):
# Build the string without the character at position i.
var prefix = String()
var suffix = String()
for j in range(0, i):
prefix.append(codepoints[j])
for k in range(i + 1, length):
suffix.append(codepoints[k])
# Recursively permute the remaining characters.
var intermediate_perms = permute_recursive(prefix + suffix)
# Append the selected character to each sub-permutation.
for ref perm in intermediate_perms:
perm.append(codepoints[i])
perms.extend(intermediate_perms^)
return perms^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Approach 2 β iterative insertion
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def permute_iterative(word: String) -> List[String]:
"""Return all permutations of `word` via incremental insertion.
Start with a single permutation containing the last character.
For each remaining character, insert it at every possible position
in every existing permutation, building up the full set.
"""
var stack = [chr(Int(codepoint)) for codepoint in word.codepoints()]
if len(stack) == 0 or len(stack) == 1:
return [word]
# Seed with the last character.
var perms = [stack.pop()]
while stack:
var letter = stack.pop()
var incremental_perms: List[String] = []
for perm in perms:
var chars = [codepoint for codepoint in perm.codepoints()]
# Insert `letter` at every possible position.
for i in range(len(chars) + 1):
var temp = String()
for j in range(i):
temp.append(chars[j])
temp.append(Codepoint.ord(letter))
for k in range(i, len(chars)):
temp.append(chars[k])
incremental_perms.append(temp^)
perms = incremental_perms^
return perms^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Helpers
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def _is_permutation_of(subject: String, source: String) -> Bool:
"""Check that `subject` contains exactly the chars of `source`.
Both strings must have the same length and the same multiset of
Unicode codepoints.
"""
if subject.count_codepoints() != source.count_codepoints():
return False
var sub_counts = Dict[Int, Int]()
var src_counts = Dict[Int, Int]()
for code in subject.codepoints():
sub_counts[Int(code)] = sub_counts.get(Int(code), 0) + 1
for code in source.codepoints():
src_counts[Int(code)] = src_counts.get(Int(code), 0) + 1
return sub_counts == src_counts
def _all_permutations_are_unique(perms: List[String]) -> Bool:
"""Return `True` if every string in `perms` is distinct."""
var seen = Set[String]()
for p in perms:
if p in seen:
return False
seen.add(p)
return True
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# ββ recursive approach ββββββββββββββββββββββββββββββββββββββ
def test_rec_example_three_chars() raises:
var result = permute_recursive("abc")
assert_equal(len(result), 6)
for p in result:
assert_true(_is_permutation_of(p, "abc"))
assert_true(_all_permutations_are_unique(result))
def test_rec_single_char() raises:
var result = permute_recursive("x")
assert_equal(len(result), 1)
assert_equal(result[0], "x")
def test_rec_two_chars() raises:
var result = permute_recursive("ab")
assert_equal(len(result), 2)
assert_true(_all_permutations_are_unique(result))
for p in result:
assert_true(_is_permutation_of(p, "ab"))
def test_rec_empty_string() raises:
var result = permute_recursive("")
assert_equal(len(result), 1)
assert_equal(result[0], "")
def test_rec_four_chars() raises:
var result = permute_recursive("abcd")
assert_equal(len(result), 24)
assert_true(_all_permutations_are_unique(result))
for p in result:
assert_true(_is_permutation_of(p, "abcd"))
def test_rec_with_repeated_chars() raises:
var result = permute_recursive("aab")
# With duplicate chars, some permutations will be identical.
assert_equal(len(result), 6) # 6 variations, some may be duplicates
for p in result:
assert_true(_is_permutation_of(p, "aab"))
# ββ iterative approach ββββββββββββββββββββββββββββββββββββββ
def test_iter_example_three_chars() raises:
var result = permute_iterative("abc")
assert_equal(len(result), 6)
for p in result:
assert_true(_is_permutation_of(p, "abc"))
assert_true(_all_permutations_are_unique(result))
def test_iter_single_char() raises:
var result = permute_iterative("z")
assert_equal(len(result), 1)
assert_equal(result[0], "z")
def test_iter_two_chars() raises:
var result = permute_iterative("xy")
assert_equal(len(result), 2)
assert_true(_all_permutations_are_unique(result))
def test_iter_empty_string() raises:
var result = permute_iterative("")
assert_equal(len(result), 1)
assert_equal(result[0], "")
def test_iter_four_chars() raises:
var result = permute_iterative("efgh")
assert_equal(len(result), 24)
assert_true(_all_permutations_are_unique(result))
# ββ cross-verification ββββββββββββββββββββββββββββββββββββββ
def test_both_implementations_produce_same_set() raises:
var inputs = ["a", "ab", "abc", "xyz", "hello"]
for word in inputs:
var r1 = permute_recursive(word)
var r2 = permute_iterative(word)
assert_equal(len(r1), len(r2))
# Both must contain the same set of strings.
for p in r1:
assert_true(p in r2)
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Longest Increasing Subsequence.
Given an integer array nums, find the length of the longest strictly
increasing subsequence (LIS), and return the subsequence itself.
Two functions are provided:
β’ length_of_lis β O(nΒ²) DP, returns only the length.
β’ longest_increasing_subsequence β same DP but also reconstructs
the actual sequence by tracking the next index in the chain.
Algorithm β O(nΒ²) time, O(n) space:
dp[i] = length of the longest strictly increasing subsequence
starting at index i. Process indices from right to left; for
each j > i where nums[i] < nums[j], update
dp[i] = max(dp[i], 1 + dp[j]).
The sequence-reconstruction variant maintains a next_idx[i] array:
when 1 + dp[j] > dp[i], set next_idx[i] = j. After the DP pass,
find the index with the maximum dp value and walk forward through
the next_idx pointers.
Example:
length_of_lis([10, 9, 2, 5, 3, 7, 101, 18]) β 4
longest_increasing_subsequence(...) β [2, 5, 7, 101]
from std.testing import assert_equal, assert_true, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Length-only (standard DP)
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def length_of_lis(nums: List[Int]) -> Int:
"""Length of the longest strictly increasing subsequence.
DP scanning right-to-left: `dp[i]` = LIS length starting at `i`.
Base: every element alone forms a subsequence of length 1.
"""
var n = len(nums)
if n == 0 or n == 1:
return n
var dp = List[Int](length=n, fill=1)
for i in range(n - 1, -1, -1):
for j in range(i + 1, n):
if nums[i] < nums[j]:
dp[i] = max(dp[i], 1 + dp[j])
var max_len = dp[0]
for i in range(1, n):
max_len = max(max_len, dp[i])
return max_len
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Sequence reconstruction
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def longest_increasing_subsequence(nums: List[Int]) -> List[Int]:
"""Return the longest strictly increasing subsequence itself.
Uses the same right-to-left DP as `length_of_lis`, but additionally
tracks the next index in the optimal chain so the sequence can be
reconstructed by a forward walk.
"""
var n = len(nums)
if n == 0:
return List[Int]()
if n == 1:
return nums.copy()
var dp = List[Int](length=n, fill=1)
var next_idx = List[Int](length=n, fill=-1)
for i in range(n - 1, -1, -1):
for j in range(i + 1, n):
if nums[i] < nums[j] and 1 + dp[j] > dp[i]:
dp[i] = 1 + dp[j]
next_idx[i] = j
# Find the starting index of the longest sequence.
var start = 0
for i in range(1, n):
if dp[i] > dp[start]:
start = i
# Walk forward through next_idx pointers to reconstruct.
var result = List[Int](capacity=dp[start])
var curr = start
while curr >= 0:
result.append(nums[curr])
curr = next_idx[curr]
return result^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Helpers
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def _is_strictly_incr(seq: List[Int]) -> Bool:
for i in range(1, len(seq)):
if seq[i - 1] >= seq[i]:
return False
return True
def _is_subsequence_of(sub: List[Int], sup: List[Int]) -> Bool:
"""Check that `sub` appears in order within `sup`."""
var si = 0
for x in sup:
if si < len(sub) and sub[si] == x:
si += 1
return si == len(sub)
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# ββ length_of_lis βββββββββββββββββββββββββββββββββββββββββββ
def test_len_example_1() raises:
assert_equal(length_of_lis([10, 9, 2, 5, 3, 7, 101, 18]), 4)
def test_len_example_2() raises:
assert_equal(length_of_lis([0, 1, 0, 3, 2, 3]), 4)
def test_len_example_3() raises:
assert_equal(length_of_lis([7, 7, 7, 7, 7, 7, 7]), 1)
def test_len_empty() raises:
assert_equal(length_of_lis([]), 0)
def test_len_single() raises:
assert_equal(length_of_lis([5]), 1)
def test_len_decreasing() raises:
assert_equal(length_of_lis([5, 4, 3, 2, 1]), 1)
def test_len_negative() raises:
assert_equal(length_of_lis([-2, -1, 0, 5]), 4)
# ββ longest_increasing_subsequence ββββββββββββββββββββββββββ
def test_seq_example_1() raises:
var result = longest_increasing_subsequence([10, 9, 2, 5, 3, 7, 101, 18])
assert_equal(len(result), 4)
assert_true(_is_strictly_incr(result))
assert_true(_is_subsequence_of(result, [10, 9, 2, 5, 3, 7, 101, 18]))
def test_seq_example_2() raises:
var result = longest_increasing_subsequence([0, 1, 0, 3, 2, 3])
assert_equal(len(result), 4)
assert_true(_is_strictly_incr(result))
assert_true(_is_subsequence_of(result, [0, 1, 0, 3, 2, 3]))
def test_seq_example_3() raises:
var result = longest_increasing_subsequence([7, 7, 7, 7, 7, 7, 7])
assert_equal(len(result), 1)
assert_true(_is_strictly_incr(result))
assert_true(_is_subsequence_of(result, [7, 7, 7, 7, 7, 7, 7]))
def test_seq_empty() raises:
assert_equal(len(longest_increasing_subsequence([])), 0)
def test_seq_single() raises:
assert_equal(longest_increasing_subsequence([42]), [42])
def test_seq_decreasing() raises:
var result = longest_increasing_subsequence([5, 4, 3, 2, 1])
assert_equal(len(result), 1)
assert_true(_is_subsequence_of(result, [5, 4, 3, 2, 1]))
def test_seq_negative() raises:
var result = longest_increasing_subsequence([-2, -1, 0, 5])
assert_equal(result, [-2, -1, 0, 5])
def test_seq_simple_ascending() raises:
assert_equal(longest_increasing_subsequence([1, 2, 3, 4]), [1, 2, 3, 4])
def test_seq_matches_length_of_lis() raises:
var case0: List[Int] = [10, 9, 2, 5, 3, 7, 101, 18]
assert_equal(length_of_lis(case0), len(longest_increasing_subsequence(case0)))
var case1: List[Int] = [0, 1, 0, 3, 2, 3]
assert_equal(length_of_lis(case1), len(longest_increasing_subsequence(case1)))
var case2: List[Int] = [7, 7, 7, 7, 7, 7, 7]
assert_equal(length_of_lis(case2), len(longest_increasing_subsequence(case2)))
var case3 = List[Int]()
assert_equal(length_of_lis(case3), len(longest_increasing_subsequence(case3)))
var case4: List[Int] = [5]
assert_equal(length_of_lis(case4), len(longest_increasing_subsequence(case4)))
var case5: List[Int] = [5, 4, 3, 2, 1]
assert_equal(length_of_lis(case5), len(longest_increasing_subsequence(case5)))
var case6: List[Int] = [-2, -1, 0, 5]
assert_equal(length_of_lis(case6), len(longest_increasing_subsequence(case6)))
var case7: List[Int] = [1, 2, 3, 4]
assert_equal(length_of_lis(case7), len(longest_increasing_subsequence(case7)))
var case8: List[Int] = [3, 1, 4, 1, 5, 9, 2, 6, 5]
assert_equal(length_of_lis(case8), len(longest_increasing_subsequence(case8)))
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Decode String.
Given an encoded string, return its decoded string. The encoding rule
is k[encoded_string] where the content inside brackets is repeated
k times. Input is always well-formed; digits appear only as repeat
counts. Nested encoding is supported.
Two implementations are provided β both O(n Β· k):
β’ decode_str β character stack. Push codepoints until ];
unwind the innermost bracket, expand, and push the result
back. Handles nesting naturally because expanded content is
available before the next ] is processed.
β’ decode_str_two_stack β count + string stacks. Walk the
input with a byte pointer; push (current_string, count) onto
stacks at [; pop and append count times at ]. Avoids
unwinding character-by-character from the stack.
Example:
decode_str("3[a]2[bc]") β "aaabcbc"
decode_str("3[a2[c]]") β "accaccacc"
decode_str("2[abc]3[cd]ef") β "abcabccdcdcdef"
from std.collections.string import Codepoint
from std.testing import assert_equal, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Approach 1 β character stack
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def decode_str(s: String) raises -> String:
"""Decode using a single codepoint stack.
Push characters until `]`. On `]`, pop back to `[` to recover
the section, read the repeat count from the digits before `[`,
expand the section, and push the expanded result back onto the
stack. After the full pass the stack holds the decoded string.
"""
var chars = s.codepoints()
if len(chars) == 0 or len(chars) == 1:
return s
var stack = List[Codepoint](capacity=len(chars))
for ch in chars:
if ch == Codepoint.ord("]"):
# ββ 1. Recover the section inside the brackets ββ
var section = String()
while len(stack) > 0 and stack[len(stack) - 1] != Codepoint.ord("["):
section = chr(Int(stack.pop())) + section
# ββ 2. Pop the opening bracket ββ
if len(stack) > 0:
_ = stack.pop()
# ββ 3. Read the repeat count (digits before `[`) ββ
var count_str = String()
while len(stack) > 0 and stack[len(stack) - 1].is_ascii_digit():
count_str = chr(Int(stack.pop())) + count_str
# ββ 4. Expand and push back ββ
var times = Int(count_str)
var expanded = String()
for _ in range(times):
expanded += section
for code in expanded.codepoints():
stack.append(code)
else:
stack.append(ch)
var result = String()
for code in stack:
result.append(code)
return result
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Approach 2 β count + string stacks
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def decode_str_two_stack(s: String) raises -> String:
"""Decode using count and string stacks (LeetCode 394 style).
Walk the input byte-by-byte:
β’ digit β accumulate the full multi-digit number and push it.
β’ `[` β push the current string onto the stack; reset.
β’ `]` β pop count and previous string; append current count
times; result becomes the new current string.
β’ else β append the byte as a codepoint to the current string.
"""
var bytes = s.as_bytes()
var n = len(bytes)
if n == 0 or n == 1:
return s
var count_stack = List[Int]()
var str_stack = List[String]()
var current = String()
var i = 0
comptime ascii_zero = UInt8(48)
while i < n:
var b = bytes[i]
if ascii_zero <= b <= UInt8(57): # '0' β¦ '9'
var num = 0
while i < n and ascii_zero <= bytes[i] <= UInt8(57):
num = num * 10 + Int(bytes[i]) - Int(ascii_zero)
i += 1
count_stack.append(num)
elif b == UInt8(Int(Codepoint.ord("["))):
str_stack.append(current^)
current = String()
i += 1
elif b == UInt8(Int(Codepoint.ord("]"))):
var times = count_stack.pop()
var prev = str_stack.pop()
for _ in range(times):
prev += current
current = prev^
i += 1
else:
current.append(Codepoint(b))
i += 1
return current
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# ββ character-stack approach ββββββββββββββββββββββββββββββββ
def test_cs_simple_repeat() raises:
assert_equal(decode_str("3[a]"), "aaa")
def test_cs_example_1() raises:
assert_equal(decode_str("3[a]2[bc]"), "aaabcbc")
def test_cs_example_2_nested() raises:
assert_equal(decode_str("3[a2[c]]"), "accaccacc")
def test_cs_example_3() raises:
assert_equal(decode_str("2[abc]3[cd]ef"), "abcabccdcdcdef")
def test_cs_single_char_no_repeat() raises:
assert_equal(decode_str("a"), "a")
def test_cs_empty_string() raises:
assert_equal(decode_str(""), "")
def test_cs_no_brackets() raises:
assert_equal(decode_str("abcdef"), "abcdef")
def test_cs_repeat_single_char() raises:
assert_equal(decode_str("5[x]"), "xxxxx")
def test_cs_double_nesting() raises:
assert_equal(decode_str("2[3[a]]"), "aaaaaa")
def test_cs_multiple_groups() raises:
assert_equal(decode_str("2[a]3[b]4[c]"), "aabbbcccc")
def test_cs_trailing_literal() raises:
assert_equal(decode_str("2[ab]c"), "ababc")
def test_cs_leading_literal() raises:
assert_equal(decode_str("x4[y]z"), "xyyyyz")
def test_cs_nested_with_literals() raises:
assert_equal(decode_str("2[a3[b]c]"), "abbbcabbbc")
# ββ two-stack approach ββββββββββββββββββββββββββββββββββββββ
def test_ts_example_1() raises:
assert_equal(decode_str_two_stack("3[a]2[bc]"), "aaabcbc")
def test_ts_example_2_nested() raises:
assert_equal(decode_str_two_stack("3[a2[c]]"), "accaccacc")
def test_ts_example_3() raises:
assert_equal(decode_str_two_stack("2[abc]3[cd]ef"), "abcabccdcdcdef")
def test_ts_empty_string() raises:
assert_equal(decode_str_two_stack(""), "")
def test_ts_single_char() raises:
assert_equal(decode_str_two_stack("a"), "a")
def test_ts_repeat() raises:
assert_equal(decode_str_two_stack("5[x]"), "xxxxx")
def test_ts_nested() raises:
assert_equal(decode_str_two_stack("2[3[a]]"), "aaaaaa")
def test_ts_multi_digit_count() raises:
assert_equal(decode_str_two_stack("12[a]"), "aaaaaaaaaaaa")
def test_ts_large_expansion() raises:
assert_equal(decode_str_two_stack("3[ab]"), "ababab")
def test_ts_leading_trailing_literals() raises:
assert_equal(decode_str_two_stack("x4[y]z"), "xyyyyz")
# ββ cross-verification ββββββββββββββββββββββββββββββββββββββ
def test_both_implementations_agree() raises:
var inputs = List[String]()
inputs.append(String("3[a]2[bc]"))
inputs.append(String("3[a2[c]]"))
inputs.append(String("2[abc]3[cd]ef"))
inputs.append(String("5[x]"))
inputs.append(String("2[3[a]]"))
inputs.append(String("x4[y]z"))
inputs.append(String("2[a3[b]c]"))
for s in inputs:
assert_equal(decode_str(s), decode_str_two_stack(s))
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Generate All Subsets.
Given an integer array nums that may contain duplicates, return all
possible subsets (the power set). The solution set must not contain
duplicate subsets.
Algorithm β O(2βΏ Β· n) time, O(2βΏ Β· n) space (output size).
Sort the array, then build subsets iteratively. Start with a result containing only the empty subset. For each element in the sorted array, extend every existing subset by appending the current element and add the new subsets to the result.
To avoid generating duplicate subsets when the input contains duplicate values, track the number of subsets that existed before processing the first occurrence of a value. When the same value appears again, only extend the subsets that were created since the previous occurrence, not the ones that already included an earlier copy.
Example:
gen_all_subsets([1, 2, 2]) β [[], [1], [1, 2], [1, 2, 2], [2], [2, 2]]
gen_all_subsets([0]) β [[], [0]]
from std.testing import assert_equal, assert_false, assert_true, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Implementation
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def gen_all_subsets(mut nums: List[Int]) -> List[List[Int]]:
"""All subsets of `nums` (power set), with duplicate handling.
Sort first, then build iteratively. For each element, extend
every existing subset by appending the element. On duplicates,
only extend subsets created since the previous occurrence to
avoid generating identical subsets.
Trace through [1, 2, 2]
βββ¬ββββββ¬ββββββββ¬ββββββββββββ¬βββββββββββββββββββ¬βββββββββββ¬ββββββββββββββββββββββββββ
βiβnums[βstart_iβend before βSubsets extended (βNew βResult after β
β βi] βdx βloop βj range) βsubsets β β
βββΌββββββΌββββββββΌββββββββββββΌβββββββββββββββββββΌβββββββββββΌββββββββββββββββββββββββββ€
β0β1 β0 β1 β[0..0] β [] β[1] β[[], [1]] β
βββΌββββββΌββββββββΌββββββββββββΌβββββββββββββββββββΌβββββββββββΌββββββββββββββββββββββββββ€
β1β2 β0 β2 β[0..1] β [], [1] β[2], [1,2]β[[], [1], [2], [1,2]] β
βββΌββββββΌββββββββΌββββββββββββΌβββββββββββββββββββΌβββββββββββΌββββββββββββββββββββββββββ€
β2β2 βend=2 β4 β[2..3] β [2], [1, β[2,2], [1,β[[], [1], [2], [1,2], [2,β
β β β β β2] β2,2] β2], [1,2,2]] β
βββ΄ββββββ΄ββββββββ΄ββββββββββββ΄βββββββββββββββββββ΄βββββββββββ΄ββββββββββββββββββββββββββ
At i=2 the duplicate 2 is found. start_idx = end = 2 β we only extend subsets at
index β₯ 2, i.e. [2] and [1,2], the subsets that were created when the first 2 was
processed. We skip [] and [1] because extending them would produce [2] and [1,2]
again β exactly the duplicates we want to avoid.
If we didn't do this and extended all 4 existing subsets at i=2, we'd get [[], [1], [
2], [1,2], *[2]*, *[1,2]*, [2,2], [1,2,2]] β 8 entries with [2] and [1,2] appearing
twice.
"""
sort(nums)
var result = List[List[Int]](capacity=2 ** len(nums))
result.append(List[Int]()) # start with the empty subset
var end: Int = 0
for i in range(len(nums)):
# If this is a duplicate, only extend subsets that were
# created since the previous occurrence of this value.
var start_idx: Int = 0
if i > 0 and nums[i] == nums[i - 1]:
start_idx = end
end = len(result)
for j in range(start_idx, end):
var new_subset = result[j].copy()
new_subset.append(nums[i])
result.append(new_subset^)
return result^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def test_example_with_duplicates() raises:
var nums: List[Int] = [1, 2, 2]
var result = gen_all_subsets(nums)
assert_equal(len(result), 6)
assert_true([] in result)
assert_true([1] in result)
assert_true([2] in result)
assert_true([1, 2] in result)
assert_true([2, 2] in result)
assert_true([1, 2, 2] in result)
def test_example_single_element() raises:
var nums: List[Int] = [0]
var result = gen_all_subsets(nums)
assert_equal(len(result), 2)
assert_equal(result[0], [])
assert_equal(result[1], [0])
def test_empty_array() raises:
var nums = List[Int]()
var result = gen_all_subsets(nums)
assert_equal(len(result), 1)
assert_equal(result[0], [])
def test_two_unique_elements() raises:
var nums: List[Int] = [1, 3]
var result = gen_all_subsets(nums)
assert_equal(len(result), 4)
assert_true([1] in result)
assert_true([3] in result)
assert_true([1, 3] in result)
assert_true([] in result)
def test_all_duplicates() raises:
var nums: List[Int] = [2, 2, 2]
var result = gen_all_subsets(nums)
assert_equal(len(result), 4)
assert_equal(result[0], [])
assert_equal(result[1], [2])
assert_equal(result[2], [2, 2])
assert_equal(result[3], [2, 2, 2])
def test_four_elements() raises:
var nums: List[Int] = [1, 2, 3, 4]
var result = gen_all_subsets(nums)
assert_equal(len(result), 16)
def test_no_duplicate_subsets() raises:
"""All subsets in the result must be unique."""
var nums: List[Int] = [1, 2, 2]
var result = gen_all_subsets(nums)
for i in range(len(result)):
for j in range(i + 1, len(result)):
assert_false(result[i] == result[j])
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Generate All Subsets (no duplicates).
Given an integer array nums of distinct integers, return all possible
subsets (the power set).
Two implementations are provided β both O(2βΏ Β· n) time and space:
β’ gen_all_subsets β iterative extension. Start with [[]] and
for each element extend every existing subset by appending it.
β’ gen_all_subsets_recursive β recursive backtracking via a
helper function with explicit mut parameters (no capturing
nested function, avoiding a Mojo compiler crash with nested
closures).
Example:
gen_all_subsets([1, 2, 3]) β [[], [1], [2], [1, 2], [3], [1, 3], [2, 3], [1, 2, 3]]
from std.testing import assert_equal, assert_false, assert_true, TestSuite
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Approach 1 β iterative
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def gen_all_subsets(mut nums: List[Int]) -> List[List[Int]]:
"""All subsets via iterative extension (no recursion).
Sort first (though input is already distinct), then for each
element extend every existing subset by appending it.
"""
sort(nums)
var result = List[List[Int]](capacity=2 ** len(nums))
result.append(List[Int]())
for i in range(len(nums)):
var end = len(result)
for j in range(end):
var new_subset = result[j].copy()
new_subset.append(nums[i])
result.append(new_subset^)
return result^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Approach 2 β recursive (no capturing)
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def backtrack(
nums: List[Int],
start: Int,
mut subset: List[Int],
mut result: List[List[Int]],
):
"""Recursive helper β appends every node of the recursion tree.
At each call, the current `subset` is a valid subset and is
appended to `result`. Then we try adding each remaining element
and recurse.
"""
if start > len(nums):
return
result.append(subset.copy())
for i in range(start, len(nums)):
subset.append(nums[i])
backtrack(nums, i + 1, subset, result)
_ = subset.pop()
def gen_all_subsets_recursive(mut nums: List[Int]) -> List[List[Int]]:
"""All subsets via recursive backtracking.
Uses a plain helper function with `mut` parameters instead of
a `capturing` nested function to avoid Mojo compiler issues
with closures.
"""
sort(nums)
var result = List[List[Int]](capacity=2 ** len(nums))
var subset = List[Int]()
backtrack(nums, 0, subset, result)
return result^
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
# Tests
# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
def _check_power_set(result: List[List[Int]], nums: List[Int]) raises:
"""Assert that `result` is a valid power set of `nums`."""
assert_equal(len(result), 2 ** len(nums))
for p in result:
for x in p:
assert_true(x in nums)
for i in range(len(result)):
for j in range(i + 1, len(result)):
assert_false(result[i] == result[j])
# ββ iterative βββββββββββββββββββββββββββββββββββββββββββββββ
def test_iter_example_three_elements() raises:
var nums: List[Int] = [1, 2, 3]
var result = gen_all_subsets(nums)
assert_equal(len(result), 8)
_check_power_set(result, [1, 2, 3])
def test_iter_single_element() raises:
var nums: List[Int] = [5]
var result = gen_all_subsets(nums)
assert_equal(len(result), 2)
_check_power_set(result, [5])
def test_iter_empty_array() raises:
var nums = List[Int]()
var result = gen_all_subsets(nums)
assert_equal(len(result), 1)
assert_equal(result[0], [])
def test_iter_four_elements() raises:
var nums: List[Int] = [10, 20, 30, 40]
var result = gen_all_subsets(nums)
assert_equal(len(result), 16)
_check_power_set(result, nums)
# ββ recursive βββββββββββββββββββββββββββββββββββββββββββββββ
def test_rec_example_three_elements() raises:
var nums: List[Int] = [1, 2, 3]
var result = gen_all_subsets_recursive(nums)
assert_equal(len(result), 8)
_check_power_set(result, [1, 2, 3])
def test_rec_single_element() raises:
var nums: List[Int] = [7]
var result = gen_all_subsets_recursive(nums)
assert_equal(len(result), 2)
_check_power_set(result, [7])
def test_rec_empty_array() raises:
var nums = List[Int]()
var result = gen_all_subsets_recursive(nums)
assert_equal(len(result), 1)
assert_equal(result[0], [])
def test_rec_four_elements() raises:
var nums: List[Int] = [2, 4, 6, 8]
var result = gen_all_subsets_recursive(nums)
assert_equal(len(result), 16)
_check_power_set(result, nums)
# ββ cross-verification ββββββββββββββββββββββββββββββββββββββ
def test_both_agree_empty() raises:
var a = List[Int]()
var r1 = gen_all_subsets(a)
var b = List[Int]()
var r2 = gen_all_subsets_recursive(b)
assert_equal(len(r1), len(r2))
def test_both_agree_single() raises:
var a: List[Int] = [5]
var r1 = gen_all_subsets(a)
var b: List[Int] = [5]
var r2 = gen_all_subsets_recursive(b)
for x in r1:
assert_true(x in r2)
for x in r2:
assert_true(x in r1)
def test_both_agree_two() raises:
var a: List[Int] = [1, 2]
var r1 = gen_all_subsets(a)
var b: List[Int] = [1, 2]
var r2 = gen_all_subsets_recursive(b)
assert_equal(len(r1), len(r2))
def test_both_agree_three() raises:
var a: List[Int] = [1, 2, 3]
var r1 = gen_all_subsets(a)
var b: List[Int] = [1, 2, 3]
var r2 = gen_all_subsets_recursive(b)
for x in r1:
assert_true(x in r2)
for x in r2:
assert_true(x in r1)
def test_both_agree_four() raises:
var a: List[Int] = [10, 20, 30, 40]
var r1 = gen_all_subsets(a)
var b: List[Int] = [10, 20, 30, 40]
var r2 = gen_all_subsets_recursive(b)
assert_equal(len(r1), len(r2))
def main() raises:
TestSuite.discover_tests[__functions_in_module()]().run()
Count the number of set bits (Hamming weight) in the binary representation of a positive integer n.
def count_bits(mut num: Int) -> Int:
result = 0
while num:
result += num % 2
num = num >> 1
return result
def count_bits2(mut num: Int) -> Int:
result = 0
while num:
result += num & 1
num >>= 1
return result
from std.testing import assert_equal
def main() raises:
num = 11
assert_equal(3, count_bits(num))
num = 128
assert_equal(1, count_bits(num))
num = 2147483645
assert_equal(30, count_bits(num))
num = 11
assert_equal(3, count_bits2(num))
num = 128
assert_equal(1, count_bits2(num))
num = 2147483645
assert_equal(30, count_bits2(num))
Count 1s For Each Entry
Return an array where each element at index i (0 β€ i β€ n) is the number of 1βs in the binary representation of i.
def count_bits(n: Int) -> List[Int]:
result = List(length=n + 1, fill=0)
power = 1
for i in range(1, n + 1):
if i == power * 2:
power = i
result[i] = 1 + result[i - power]
return result
from std.testing import assert_equal
def main() raises:
n = 2
result = count_bits(n)
assert_equal([0, 1, 1], result, "Assertion failed")
n = 5
result = count_bits(n)
assert_equal([0, 1, 1, 2, 1, 2], result, "Assertion failed")
Game Of Life
from gridv1 import Grid as GridV1
from gridv2 import Grid as GridV2
from utils import Variant
import random
comptime Grid = Variant[GridV1, GridV2]
def run(owned grid: Grid) raises -> None:
var inner: GridV1
if grid.isa[GridV1]():
inner = grid[GridV1]
print("Received a grid of type V1")
else:
inner = GridV1(grid[GridV2])
print("Received a grid of type V2 - converted to V1")
while True:
print("Current mutation:\n\n")
print(inner)
print()
print()
if input("Enter 'q' to quit or press <Enter> to continue: ") == "q":
break
inner.mutate()
def main() raises -> None:
random.seed()
var grid: Grid
if random.random_ui64(0, 1):
v1 = GridV1.new(16, 16)
grid = Grid(v1)
else:
v2 = GridV2.new(None, 16, 16)
grid = Grid(v2)
run(grid^)
Gridv1
import random
from gridv2 import Grid as GridV2
# Grid is a 2D structure holding cell states (0: dead, 1: alive)
# It supports string conversion, output writing, and cell access/update
@value
struct Grid(Stringable, Writable):
var data: List[List[Int, True]] # 2D grid of integers (1 = alive, 0 = dead)
# Constructor to initialize the grid with given data
def __init__(out self, data: List[List[Int, True]]):
self.data = data
@implicit
def __init__(out self, source: GridV2):
data = source.data
rows = source.rows
cols = source.cols
grid = List[List[Int, True]]()
for row in range(rows):
curr_row = List[Int, True]()
for col in range(cols):
curr_row.append(Int((data + (row * cols + col))[]))
grid.append(curr_row)
self.data = grid^
# Get the number of rows in the grid
def row_count(self) -> Int:
if self.data:
return len(self.data)
else:
return 0
# Get the number of columns in the grid
def col_count(self) -> Int:
if self.data[0]:
return len(self.data[0])
else:
return 0
# Convert the grid to a string for pretty-printing
def __str__(self) -> String:
capacity = self.row_count() * self.col_count()
if capacity == 0:
return String()
s = String(capacity=capacity)
row_index = 0
for row in self.data:
for col in row[]:
if col[] == 1:
s += "*" # Alive cell represented by '*'
else:
s += " " # Dead cell is blank
if row_index != self.row_count() - 1:
s += "\n" # Line break between rows
row_index += 1
return s
# Allow writing the grid to any output writer
def write_to[W: Writer](self, mut writer: W) -> None:
writer.write(self.__str__())
# Access cell at (row, col)
def __getitem__(self, row: Int, col: Int) -> Int:
return self.data[row][col]
# Update cell at (row, col)
def __setitem__(mut self, row: Int, col: Int, value: Int) -> None:
self.data[row][col] = value
# Static method to create a random grid with specified size
@staticmethod
def new(rows: Int, cols: Int) -> Self:
random.seed()
data = List[List[Int, True]](capacity=rows)
for row in range(rows):
record = List[Int, True](capacity=cols)
for col in range(cols):
# Initialize each cell randomly to 0 or 1
record.append(Int(random.random_si64(0, 1)))
data.append(record)
return Self(data)
# Perform one step of mutation (Game of Life rules)
def mutate(mut self):
rows = self.row_count()
cols = self.col_count()
for row in range(rows):
above = (row - 1) % rows
below = (row + 1) % rows
for col in range(cols):
left = (col - 1) % cols
right = (col + 1) % cols
# Count live neighbors using 8-connected grid
alive_neighbours = (
self[above, left]
+ self[above, col]
+ self[above, right]
+ self[row, right]
+ self[below, right]
+ self[below, col]
+ self[below, left]
+ self[row, left]
)
# Apply Conway's Game of Life rules:
# Rule 1 & 2: Any live cell with 2 or 3 live neighbors survives
if self[row, col] == 1 and (
alive_neighbours == 2 or alive_neighbours == 3
):
continue # Keep alive
# Rule 3: All other live cells die
else:
self[row, col] = 0
# Rule 4: Any dead cell with exactly 3 live neighbors becomes alive
if self[row, col] == 0 and alive_neighbours == 3:
self[row, col] = 1
def run(owned grid: Grid) raises -> None:
while True:
print("Current mutation:\n\n")
print(grid)
print()
print()
if input("Enter 'q' to quit or press <Enter> to continue: ") == "q":
break
grid.mutate()
def main() raises -> None:
grid_2 = GridV2.new(42, 16, 16)
#run(grid_2)
print(grid_2)
print("Implicit conversion\n\n")
grid_1 = Grid(grid_2)
print(grid_1)
Gridv2
import random
from collections import Optional
from memory import UnsafePointer, memcpy, memset_zero
from gridv1 import Grid as GridV1
struct Grid(Stringable, Writable):
var data: UnsafePointer[UInt8]
var rows: Int
var cols: Int
def __init__(out self, rows: Int, cols: Int):
self.rows = rows
self.cols = cols
self.data = UnsafePointer[UInt8].alloc(rows * cols)
def __init__(out self, source: GridV1):
rows = len(source.data)
cols = len(source.data[0])
self = Self(rows, cols)
for row in range(rows):
for col in range(cols):
value = UInt8(source[row, col])
(self.data + row * cols + col)[] = value
def __copyinit__(out self, existing: Self):
self.rows = existing.rows
self.cols = existing.cols
count = self.rows * self.cols
self.data = UnsafePointer[UInt8].alloc(count)
memcpy(dest=self.data, src=existing.data, count=count)
def __moveinit__(out self, owned existing: Self):
self.data = existing.data
self.rows = existing.rows
self.cols = existing.cols
def __del__(owned self):
self.data.free()
def __str__(self) -> String:
capacity = self.rows * self.cols
if capacity == 0:
return String()
s = String(capacity=capacity)
for row in range(self.rows):
for col in range(self.cols):
# if (self.data + row * self.cols + col)[] == 1:
if self[row, col] == 1:
s += "*" # Alive cell represented by '*'
else:
s += " " # Dead cell is blank
if row != self.rows - 1:
s += "\n" # Line break between rows
return s
def write_to[W: Writer](self, mut writer: W) -> None:
writer.write(self.__str__())
def __getitem__(self, row: Int, col: Int) -> UInt8:
return (self.data + row * self.cols + col)[]
def __setitem__(mut self, row: Int, col: Int, value: UInt8) -> None:
(self.data + row * self.cols + col)[] = value
@staticmethod
def new(seed: Optional[Int], rows: Int, cols: Int) -> Self:
if seed:
random.seed(seed.value())
else:
random.seed()
grid = Self(rows, cols)
random.randint(grid.data, rows * cols, 0, 1)
return grid
def mutate(mut self) -> None:
rows = self.rows
cols = self.cols
for row in range(rows):
above = (row - 1) % rows
below = (row + 1) % rows
for col in range(cols):
left = (col - 1) % cols
right = (col + 1) % cols
alive_neighbours = (
self[above, left]
+ self[above, col]
+ self[above, right]
+ self[row, right]
+ self[below, right]
+ self[below, col]
+ self[below, left]
+ self[row, left]
)
if self[row, col] == 1:
if alive_neighbours < 2:
self[row, col] = 0
if alive_neighbours == 2 or alive_neighbours == 3:
continue
if alive_neighbours > 3:
self[row, col] = 0
else:
if alive_neighbours == 3:
self[row, col] = 1
def run(owned grid: Grid) raises -> None:
while True:
print("Current mutation:\n\n")
print(grid)
print()
print()
if input("Enter 'q' to quit or press <Enter> to continue: ") == "q":
break
grid.mutate()
def main() raises -> None:
grid_1 = GridV1.new(16, 16)
# run(grid_1)
print(grid_1)
print("Implicit conversion\n\n")
grid_2 = Grid(grid_1)
print(grid_2)
Cyclic Reference 1
# We have no issues Referece1 calling Reference2 which also calls Reference1
from cyclic_reference_2 import Reference2
struct Reference1:
def __init__(out self):
Reference2.print("Reference2 inside Reference1 constructor")
@staticmethod
def print(s: String):
print(s)
def main():
var ref1 = Reference1()
Cyclic Reference 2
# We have no issues Referece2 calling Reference1 which also calls Reference2
from cyclic_reference_1 import Reference1
struct Reference2:
def __init__(out self):
Reference1.print("Reference1 inside Reference2 constructor")
@staticmethod
def print(s: String):
print(s)
def main():
var ref2 = Reference2()
Buffer Reduce
from buffer import NDBuffer
from algorithm import vectorize
from sys import simdwidthof
def summer[
type: DType, //, simdwidth: Int = simdwidthof[type]()
](buffer: NDBuffer[type=type, rank=1]) -> Scalar[type]:
result = Scalar[type](0)
@parameter
def sum[simd_width: Int](idx: Int):
result += buffer.load[width=simd_width](idx).reduce_add()
vectorize[sum, simdwidth](len(buffer))
return result
from collections import InlineArray
from math import iota
def main() raises:
comptime elem_count = 30
var array = InlineArray[Scalar[DType.float64], elem_count](
uninitialized=True
)
iota(array.unsafe_ptr(), elem_count)
var buf = NDBuffer[DType.float64, 1, _, elem_count](array)
result = summer(buf)
print(result)
Custom Struct Compare
from search_sorted_rotated_arr import find
from shapes import Rectangle
def main():
# Re
r4 = Rectangle(4, 10) # 40
r5 = Rectangle(4, 12) # 48
r6 = Rectangle(5, 10) # 50
r7 = Rectangle(8, 8) # 64
r1 = Rectangle(3, 2) # 6
r2 = Rectangle(4, 4) # 16
r3 = Rectangle(4, 8) # 32
# Rectangles are sorted and rotated in the list
items = [r4, r5, r6, r7, r1, r2, r3]
item_index = find(items, r2)
debug_assert(item_index == 5, "Assertion failed")
SIMD Select
Select based on a SIMD (Single Instruction, Multiple Data) mask This example demonstrates how to use SIMD.select() to perform element-wise conditional selection between two SIMD vectors based on a boolean mask.
from std.testing import assert_true
def main() raises:
# Create a SIMD (Single Instruction, Multiple Data) boolean selector vector of size 4.
# Each element is a boolean value (True or False), indicating which value to select from `left` or `right`.
# - If selector[i] == True, take the value from `left[i]`
# - If selector[i] == False, take the value from `right[i]`
selector = SIMD[DType.bool, 4](False, True, False, True)
# Define a SIMD vector `left` with 4 elements of unsigned 8-bit integers
left = SIMD[DType.uint8, 4](0, 42, 0, 42)
# Define another SIMD vector `right` with 4 elements of unsigned 8-bit integers
right = SIMD[DType.uint8, 4](42, 0, 42, 0)
# Use the selector to choose elements from either `left` or `right`:
# result[i] = left[i] if selector[i] else right[i]
result = selector.select(left, right)
# β result = [42, 42, 42, 42]
expected = SIMD[DType.uint8, 4](42) # β expected = [42, 42, 42, 42]
assert_true(all(result == expected))
Check device core
Check physical and logical cores of the device
from sys import num_physical_cores, num_logical_cores
def main():
print(" Physical Cores : ", num_physical_cores())
print(" Logical Cores : ", num_logical_cores())
Add 10
Implement a kernel that adds 10 to each position of vector a and stores it in vector out.
from gpu.host import DeviceContext
from memory import UnsafePointer
from gpu import thread_idx
comptime SIZE = 4
comptime BLOCKS_PER_GRID = 1
comptime THREADS_PER_BLOCK = SIZE
comptime dtype = DType.float32
def add_10(
out: UnsafePointer[Scalar[dtype]], array: UnsafePointer[Scalar[dtype]]
):
tid = thread_idx.x
out[tid] = array[tid] + 10
def main() raises:
ctx = DeviceContext()
d_array_buff = ctx.enqueue_create_buffer[dtype](SIZE)
expected = ctx.enqueue_create_buffer[dtype](SIZE)
d_out_buff = ctx.enqueue_create_buffer[dtype](SIZE)
_ = d_out_buff.enqueue_fill(0)
with d_array_buff.map_to_host() as h_array_buff:
for i in range(SIZE):
h_array_buff[i] = i
ctx.enqueue_function[add_10](
d_out_buff.unsafe_ptr(),
d_array_buff.unsafe_ptr(),
grid_dim=BLOCKS_PER_GRID,
block_dim=THREADS_PER_BLOCK,
)
ctx.synchronize()
with d_out_buff.map_to_host() as h_out_buff:
print(h_out_buff)
Add a constant 10
Implement a kernel that adds 10 to each position of 2d matrix a and stores it in out 2d matrix.
from gpu.host import DeviceContext
from memory import UnsafePointer
from gpu import thread_idx, block_dim
from std.testing import assert_equal
comptime SIZE = 2
comptime BLOCKS_PER_GRID = 1
comptime THREADS_PER_BLOCK = (3,3)
comptime dtype = DType.float32
def add_10_2d(
out: UnsafePointer[Scalar[dtype]], array: UnsafePointer[Scalar[dtype]], size: Int
):
tid = thread_idx.z * (block_dim.y * block_dim.x) + thread_idx.y * block_dim.x + thread_idx.x
if tid < size * size:
out[tid] = array[tid] + 10
def main():
try:
ctx = DeviceContext()
d_array_buff = ctx.enqueue_create_buffer[dtype](SIZE * SIZE).enqueue_fill(0)
d_out_buff = ctx.enqueue_create_buffer[dtype](SIZE * SIZE).enqueue_fill(0)
expected = ctx.enqueue_create_host_buffer[dtype](SIZE * SIZE).enqueue_fill(0)
with d_array_buff.map_to_host() as h_array_buff:
for i in range(SIZE):
for j in range(SIZE):
h_array_buff[i * SIZE + j] = i * SIZE + j
expected[i * SIZE + j] = h_array_buff[i * SIZE + j] + 10
print("Input: ", h_array_buff)
ctx.enqueue_function[add_10_2d](
d_out_buff.unsafe_ptr(),
d_array_buff.unsafe_ptr(),
SIZE,
grid_dim=BLOCKS_PER_GRID,
block_dim=THREADS_PER_BLOCK,
)
ctx.synchronize()
with d_out_buff.map_to_host() as h_out_buff:
print(h_out_buff)
print(expected)
for i in range(SIZE * SIZE ):
assert_equal(h_out_buff[i], expected[i])
except e:
print(e)
Add constant to 2D Layout tensor
Implement a kernel that adds 10 to each position of 2D LayoutTensor a and stores it in 2D LayoutTensor out.
from gpu.host import DeviceContext
from gpu import thread_idx
from layout import Layout, LayoutTensor
from math import iota
comptime SIZE = 2
comptime BLOCKS_PER_GRID = 1
comptime THREADS_PER_BLOCK = (3, 3)
comptime dtype = DType.float32
comptime layout = Layout.row_major(SIZE, SIZE)
def add_10_2dlayout(
out: LayoutTensor[mut=True, dtype, layout],
a: LayoutTensor[mut=True, dtype, layout],
size: Int,
):
row = thread_idx.y
col = thread_idx.x
# FILL ME IN (roughly 2 lines)
if row < size and col < size:
out[row, col] = a[row, col] + 10
def main():
try:
ctx = DeviceContext()
buffer_a = ctx.enqueue_create_buffer[dtype](SIZE * SIZE).enqueue_fill(
0.0
)
buffer_out = ctx.enqueue_create_buffer[dtype](SIZE * SIZE).enqueue_fill(
0.0
)
with buffer_a.map_to_host() as h_buffer_a:
iota(h_buffer_a.unsafe_ptr(), SIZE * SIZE)
out = LayoutTensor[mut=True, dtype, layout](buffer_out)
a = LayoutTensor[mut=True, dtype, layout](buffer_a)
ctx.enqueue_function[add_10_2dlayout](
out,
a,
SIZE,
grid_dim=(BLOCKS_PER_GRID, BLOCKS_PER_GRID),
block_dim=THREADS_PER_BLOCK,
)
ctx.synchronize()
with buffer_out.map_to_host() as h_buffer_out:
print(h_buffer_out)
except e:
print(e)
Add 10
Implement a kernel that adds 10 to each position of vector a and stores it in vector out. More threads than data β guard against out-of-bounds access.
from gpu.host import DeviceContext
from memory import UnsafePointer
from gpu import thread_idx, block_dim, block_idx
from std.testing import assert_equal
comptime SIZE = 4
comptime BLOCKS_PER_GRID = 1
comptime THREADS_PER_BLOCK = (8, 1)
comptime dtype = DType.float32
def add_10_with_guard(
out: UnsafePointer[Scalar[dtype]], array: UnsafePointer[Scalar[dtype]]
):
tid = (
thread_idx.z * (block_dim.y * block_dim.x)
+ thread_idx.y * block_dim.x
+ thread_idx.x
)
if tid < SIZE:
out[tid] = array[tid] + 10
def main() raises:
ctx = DeviceContext()
d_array_buff = ctx.enqueue_create_buffer[dtype](SIZE)
d_out_buff = ctx.enqueue_create_buffer[dtype](SIZE)
expected = ctx.enqueue_create_host_buffer[dtype](SIZE)
_ = d_out_buff.enqueue_fill(0)
with d_array_buff.map_to_host() as h_array_buff:
for i in range(SIZE):
h_array_buff[i] = i
ctx.enqueue_function[add_10_with_guard](
d_out_buff.unsafe_ptr(),
d_array_buff.unsafe_ptr(),
grid_dim=BLOCKS_PER_GRID,
block_dim=THREADS_PER_BLOCK,
)
ctx.synchronize()
for i in range(SIZE):
expected[i] = i + 10
print(expected)
with d_out_buff.map_to_host() as h_out_buff:
print(h_out_buff)
for i in range(SIZE):
assert_equal(h_out_buff[i], expected[i])
Vectorized Memory Access & SIMD Grid-Stride Add.
Most GPU kernels are bandwidth bound: the bottleneck is moving bytes between global
memory and the ALUs, not the arithmetic itself. A scalar load/store instruction
(LDG.E / STG.E in CUDA) moves 32 bits per operation. Modern GPUs also support
wider instructions β LDG.E.64, LDG.E.128 β that transfer 64 or 128 bits in a
single op, cutting instruction count by 2-4Γ and improving throughput.
The NVIDIA blog post CUDA Pro Tip: Increase Performance with Vectorized Memory
Access lays out the basic technique in CUDA C++: cast an int* to int2* or
int4* and the compiler generates the wider loads. A scalar copy loop that
processes one element per thread becomes 2Γ (int2) or 4Γ (int4) fewer instructions,
directly raising bandwidth utilisation. The key requirement is alignment β device
memory is naturally aligned to the vector width.
This Mojo kernel goes one step further.
βββ What this kernel does βββ
β’ Uses Mojoβs native SIMD intrinsics (load[width=N] / store[width=N]) for
vectorised memory access. The SIMD width is auto-detected from the data type
(simd_width_of[dtype]()), so the code adapts to float32 (width 4), float64
(width 2), etc.
β’ Each thread processes simd_vectors_per_thread Γ simd_width elements per grid
step β a compile-time unrolled block of SIMD vectors. The blog post stops at
one vectorised element per thread per iteration; here a single thread moves
16 float32 values at once (e.g. 4 vectors Γ width 4).
β’ A grid-stride outer loop makes the kernel grid-size-agnostic: each thread
strides across the array by grid_dim Γ block_dim, so the same kernel works
whether you launch 16 blocks or 1600.
β’ A scalar tail loop handles leftover elements that donβt fill a full SIMD vector, without warping the fast path.
β’ Runs on both CPU (DeviceContext(api="cpu")) and GPU (when an accelerator
is available), comparing results with assert_almost_equal.
β’ A companion CUDA implementation lives in gpu/vector_add.cu, showing the
same ideas β float4 vectorised loads and #pragma unroll β in native CUDA C++.
βββ Data-flow diagram βββ
Global memory: [ e0 β e1 β e2 β e3 β e4 β e5 β e6 β e7 β ... β eN ]
β
βββββββββββββββ΄βββββββββββββββββββ
β Grid-stride loop (outer) β
β thread t starts at β
β t Γ CHUNK_SIZE and advances β
β by grid_dim Γ block_dim chunksβ
βββββββββββββββ¬βββββββββββββββββββ
β
βββββββββββββββ΄ββββββββββββββββββββββββββ
β CHUNK_SIZE per iteration β
β βββββ simd_vectors_per_thread βββββ β
β β vector 0 β vector 1 β ... β β
β β ββββββββ ββββββββ ββββββββ β β
β β β SIMD β β SIMD β β SIMD β β β
β β β WWWW β β WWWW β β WWWW β β β
β β ββββββββ ββββββββ ββββββββ β β
β βββββββββββββββββββββββββββββββββββ β
β Each SIMD block = simd_width elementsβ
β loaded/stored as one unit β
βββββββββββββββ¬ββββββββββββββββββββββββββ
β
βββββββββββββββ΄βββββββββββββββββββ
β Wide load / store: β
β a.load[width=4](i) β
β β LDG.E.128 (4 Γ float32) β
β β
β result.store[width=4](i, ...) β
β β STG.E.128 (4 Γ float32) β
ββββββββββββββββββββββββββββββββββ
βββ Running βββ
pixi run mojo -I . gpu/vector_add.mojo
The kernel fills two random vectors (or uses a seed for reproducibility), adds them on CPU for reference, then on GPU (if one is available), and asserts element-wise equality within a tolerance.
from std.gpu.host import DeviceContext, HostBuffer, DeviceAttribute
from std.gpu import thread_idx, block_idx, block_dim, grid_dim
from std.testing import assert_almost_equal
from utils import Timer
from std.random import random_float64, seed
from std.sys import has_accelerator, simd_width_of
# GPU kernel: element-wise vector addition with grid-stride loop, SIMD loads,
# and compile-time loop unrolling. Each thread processes CHUNK_SIZE elements
# per iteration, then advances by the total grid stride.
#
# Parameters:
# result: output pointer (mutably addressed)
# a, b: input pointers (immutably addressed)
# size: number of elements in each vector
#
# Template parameters:
# dtype: element data type (e.g. DType.float32)
# simd_width: SIMD width, auto-detected from dtype
# simd_vectors_per_thread: number of SIMD vectors per thread per grid step
#
def vector_add[
dtype: DType,
simd_width: Int = simd_width_of[dtype](),
simd_vectors_per_thread: Int = 4 * simd_width,
](
result: UnsafePointer[Scalar[dtype], MutAnyOrigin],
a: UnsafePointer[Scalar[dtype], ImmutAnyOrigin],
b: UnsafePointer[Scalar[dtype], ImmutAnyOrigin],
size: Int,
):
var tid = block_idx.x * block_dim.x + thread_idx.x
var grid_stride = grid_dim.x * block_dim.x
comptime CHUNK_SIZE = simd_vectors_per_thread * simd_width
# =========================================================
# Each thread processes CHUNK_SIZE elements
# =========================================================
var start_index = (
tid * CHUNK_SIZE
) # Start index for each thread per grid_stride
while start_index < size:
comptime for vector in range(simd_vectors_per_thread):
var i = start_index + vector * simd_width
# Bound check for this vector
if i + simd_width <= size:
# Load whole vectors, add up and store
result.store[width=simd_width](
i, a.load[width=simd_width](i) + b.load[width=simd_width](i)
)
else: # i < size, can not load a simd_length vector, handle tail
for j in range(i, size):
result.store[width=1](
j, a.load[width=1](j) + b.load[width=1](j)
)
start_index += grid_stride * CHUNK_SIZE
# CPU reference implementation: simple sequential element-wise vector addition.
#
# Parameters:
# result: output host buffer
# a, b: input host buffers
# size: number of elements
#
def vector_add_cpu[
dtype: DType,
//,
](
result: HostBuffer[dtype],
a: HostBuffer[dtype],
b: HostBuffer[dtype],
size: Int,
):
var i = 0
while i < size:
result[i] = a[i] + b[i]
i += 1
# Fill a host buffer with random float64 values cast to the target dtype.
# Optionally accepts a seed for reproducible results.
#
# Parameters:
# buffer_a: host buffer to fill
# init_seed: optional RNG seed (deterministic if provided)
# min, max: range for random values
#
def fill[
dtype: DType,
//,
](
buffer_a: HostBuffer[dtype],
init_seed: Optional[Int] = None,
min: Float64 = 1.0,
max: Float64 = 10.0,
):
if init_seed:
seed(init_seed.value())
else:
seed()
for i in range(len(buffer_a)):
buffer_a[i] = random_float64(min, max).cast[dtype]()
# Benchmark vector addition on CPU and (if available) GPU, then validate that
# all GPU results match the CPU reference within a small tolerance.
#
def main() raises:
comptime dtype = DType.float32
var size = 100000000
var cpu_ctx = DeviceContext(api="cpu")
var lhs_host_buffer = cpu_ctx.enqueue_create_host_buffer[dtype](size)
var rhs_host_buffer = cpu_ctx.enqueue_create_host_buffer[dtype](size)
var result_host_buffer = cpu_ctx.enqueue_create_host_buffer[dtype](size)
fill(lhs_host_buffer, init_seed=42)
fill(rhs_host_buffer, init_seed=123)
with Timer("CPU execution took: "):
vector_add_cpu(
result_host_buffer, lhs_host_buffer, rhs_host_buffer, size
)
cpu_ctx.synchronize()
comptime if has_accelerator():
var gpu_ctx = DeviceContext()
var result_gpu_buffer = gpu_ctx.enqueue_create_buffer[dtype](size)
var lhs_gpu_buffer = gpu_ctx.enqueue_create_buffer[dtype](size)
var rhs_gpu_buffer = gpu_ctx.enqueue_create_buffer[dtype](size)
lhs_host_buffer.enqueue_copy_to(dst=lhs_gpu_buffer)
rhs_host_buffer.enqueue_copy_to(dst=rhs_gpu_buffer)
var max_blocks_per_sm = gpu_ctx.get_attribute(
DeviceAttribute.MAX_BLOCKS_PER_MULTIPROCESSOR
)
var sm_count = gpu_ctx.get_attribute(
DeviceAttribute.MULTIPROCESSOR_COUNT
)
var threads_per_block = 256
var max_threads_per_sm = gpu_ctx.get_attribute(
DeviceAttribute.MAX_THREADS_PER_MULTIPROCESSOR
)
var max_blocks = max_threads_per_sm // threads_per_block
var blocks_count = min(max_blocks_per_sm, max_blocks) * sm_count * 4
print("Max block per sm: ", max_blocks, "sm count: ", sm_count)
print(
"Launching",
blocks_count,
"blocks with",
threads_per_block,
"threads per block",
)
with Timer("GPU execution took: "):
gpu_ctx.enqueue_function[vector_add[dtype]](
result_gpu_buffer.unsafe_ptr(),
lhs_gpu_buffer.unsafe_ptr(),
rhs_gpu_buffer.unsafe_ptr(),
size,
grid_dim=blocks_count,
block_dim=threads_per_block,
)
gpu_ctx.synchronize()
with result_gpu_buffer.map_to_host() as gpu_result:
for i in range(size):
assert_almost_equal(gpu_result[i], result_host_buffer[i])
Layout Basics
from std.gpu.host import DeviceContext
from std.sys import has_accelerator
from layout import Layout, LayoutTensor
comptime HEIGHT = 2
comptime WIDTH = 3
comptime dtype = DType.float32
comptime layout = Layout.row_major(HEIGHT, WIDTH)
comptime BLOCKS_PER_GRID = 1
comptime THREADS_PER_BLOCK = 1
def kernel[
dtype: DType, layout: Layout
](data: UnsafePointer[Scalar[dtype], MutAnyOrigin]):
var tensor = LayoutTensor[mut=True, dtype, layout, _](data)
print("Before\n")
print(tensor)
tensor[0, 0] += 1.0
print()
print("After\n")
print(tensor)
def main() raises:
var host_buffer = DeviceContext(api="cpu").enqueue_create_host_buffer[
dtype
](HEIGHT * WIDTH)
for i in range(HEIGHT * WIDTH):
host_buffer[i] = Float32(i**2)
comptime if has_accelerator():
var ctx = DeviceContext()
var device_buffer = ctx.enqueue_create_buffer[dtype](HEIGHT * WIDTH)
device_buffer.enqueue_fill(0)
host_buffer.enqueue_copy_to(device_buffer)
ctx.enqueue_function[kernel[dtype, layout]](
device_buffer.unsafe_ptr(),
grid_dim=BLOCKS_PER_GRID,
block_dim=THREADS_PER_BLOCK,
)
ctx.synchronize()
else:
var cpu_buffer = DeviceContext(api="cpu").enqueue_create_buffer[dtype](
HEIGHT * WIDTH
)
cpu_buffer.enqueue_fill(0)
host_buffer.enqueue_copy_to(cpu_buffer)
kernel[dtype, layout](cpu_buffer.unsafe_ptr())
print(host_buffer)
Dumb matrix multiplication
Simulate the CPU-style triple for-loop truly dumb matrix multiplication
from gpu.host import DeviceContext, HostBuffer
from gpu import thread_idx, block_idx, block_dim
import random
from layout import Layout, LayoutTensor
from memory import UnsafePointer, memcpy
from python import Python, PythonObject
from std.testing import assert_true
comptime ROWS_A = 8
comptime COLS_A = 16
comptime ROWS_B = 16
comptime COLS_B = 8
comptime ROWS_C = 8
comptime COLS_C = 8
comptime MATRIX_MIN_ELEM = -5.0
comptime MATRIX_MAX_ELEM = 5.0
comptime dtype = DType.float32
# Num threads per block
comptime THREADS = 1
# Total numbers blocks in the grid
comptime BLOCKS = 1
comptime layout_a = Layout.row_major(ROWS_A, COLS_A)
comptime layout_b = Layout.row_major(ROWS_B, COLS_B)
comptime layout_c = Layout.row_major(ROWS_C, COLS_C)
# alias Matrix = LayoutTensor[dtype, _, MutableAnyOrigin]
comptime Matrix = LayoutTensor[mut=True, dtype, _]
def naive_matmaul(
A: UnsafePointer[Scalar[dtype]],
B: UnsafePointer[Scalar[dtype]],
C: UnsafePointer[Scalar[dtype]],
):
var tid = block_idx.x * block_dim.x + thread_idx.x
if tid == 0:
for i in range(ROWS_A):
for j in range(COLS_B):
for k in range(COLS_A):
(C + i * COLS_C + j)[] += (A + i * COLS_A + k)[] * (
B + k * COLS_B + j
)[]
# Initialize the matrix buffer with values in the range 0 to 100
def fill_buffer(buffer: HostBuffer[dtype]):
# Randomize
# random.seed()
for i in range(len(buffer)):
buffer[i] = random.random_float64(
MATRIX_MIN_ELEM, MATRIX_MAX_ELEM
).cast[dtype]()[0]
def main():
try:
ctx = DeviceContext()
buffer_a = ctx.enqueue_create_buffer[dtype](
ROWS_A * COLS_A
).enqueue_fill(0.0)
buffer_b = ctx.enqueue_create_buffer[dtype](
ROWS_B * COLS_B
).enqueue_fill(0.0)
buffer_c = ctx.enqueue_create_buffer[dtype](
ROWS_C * COLS_C
).enqueue_fill(0.0)
with buffer_a.map_to_host() as h_buffer_a:
fill_buffer(h_buffer_a)
with buffer_b.map_to_host() as h_buffer_b:
fill_buffer(h_buffer_b)
# matrix_a = LayoutTensor[dtype, layout_a, MutableAnyOrigin](buffer_a)
# matrix_b = LayoutTensor[dtype, layout_b, MutableAnyOrigin](buffer_b)
# matrix_c = LayoutTensor[dtype, layout_c, MutableAnyOrigin](buffer_c)
ctx.enqueue_function[naive_matmaul](
buffer_a.unsafe_ptr(),
buffer_b.unsafe_ptr(),
buffer_c.unsafe_ptr(),
grid_dim=BLOCKS,
block_dim=THREADS,
)
ctx.synchronize()
with buffer_a.map_to_host() as h_buffer_a:
with buffer_b.map_to_host() as h_buffer_b:
with buffer_c.map_to_host() as h_buffer_c:
assert_allclose(
(ROWS_A, COLS_A, h_buffer_a),
(ROWS_B, COLS_B, h_buffer_b),
(ROWS_C, COLS_C, h_buffer_c),
)
except e:
print("Prininting here: ", e)
def assert_allclose(
buff_a_with_dims: (Int, Int, HostBuffer[dtype]),
buff_b_with_dims: (Int, Int, HostBuffer[dtype]),
buff_c_with_dims: (Int, Int, HostBuffer[dtype]),
) raises:
a_rows, a_cols, a_buff = buff_a_with_dims
matrix_a = reshape(to_ndarray(a_buff), a_rows, a_cols)
b_rows, b_cols, b_buff = buff_b_with_dims
matrix_b = reshape(to_ndarray(b_buff), b_rows, b_cols)
c_rows, c_cols, c_buff = buff_c_with_dims
matrix_c = reshape(to_ndarray(c_buff), c_rows, c_cols)
np = Python.import_module("numpy")
assert_true(np.allclose(np.matmul(matrix_a, matrix_b), matrix_c))
print("Assertion was successful")
def to_ndarray(buffer: HostBuffer[dtype]) raises -> PythonObject:
np = Python.import_module("numpy")
ndarray = np.zeros(len(buffer), dtype=np.float32)
ndarray_ptr = ndarray_ptr[dtype](ndarray)
buffer_ptr = buffer.unsafe_ptr()
memcpy(ndarray_ptr, buffer_ptr, len(buffer))
return ndarray
def reshape(ndarray: PythonObject, rows: Int, cols: Int) raises -> PythonObject:
return ndarray.reshape(rows, cols)
def ndarray_ptr[
dtype: DType
](ndarray: PythonObject) raises -> UnsafePointer[Scalar[dtype]]:
return ndarray.__array_interface__["data"][0].unsafe_get_as_pointer[dtype]()
Matrix multiplication 1 GPU thread per output column
Simulate the CPU-style dumb matrix multiplication 1 thread per output column
from gpu.host import DeviceContext, HostBuffer
from gpu import thread_idx, block_idx, block_dim
import random
from layout import Layout, LayoutTensor
from memory import UnsafePointer, memcpy
from python import Python, PythonObject
from std.testing import assert_true
comptime ROWS_A = 33
comptime COLS_A = 13
comptime ROWS_B = 13
comptime COLS_B = 8
comptime ROWS_C = ROWS_A
comptime COLS_C = COLS_B
comptime MATRIX_MIN_ELEM = -5.0
comptime MATRIX_MAX_ELEM = 5.0
comptime dtype = DType.float32
# Num threads per block
comptime THREADS = COLS_C
# Total numbers blocks in the grid
comptime BLOCKS = 1
comptime layout_a = Layout.row_major(ROWS_A, COLS_A)
comptime layout_b = Layout.row_major(ROWS_B, COLS_B)
comptime layout_c = Layout.row_major(ROWS_C, COLS_C)
comptime MatrixA = LayoutTensor[dtype, layout_a, MutableAnyOrigin]
comptime MatrixB = LayoutTensor[dtype, layout_b, MutableAnyOrigin]
comptime MatrixC = LayoutTensor[dtype, layout_c, MutableAnyOrigin]
def naive_matmul_one_thread_per_col[
a: Layout, b: Layout, c: Layout
](A: MatrixA, B: MatrixB, C: MatrixC,):
var tid = block_idx.x * block_dim.x + thread_idx.x
if tid < COLS_C: # Each thread id `tid` is cols of C or B
for i in range(ROWS_A):
for k in range(COLS_A):
C[i, tid] += A[i, k] * B[k, tid]
# Initialize the matrix buffer with values in the range 0 to 100
def fill_buffer(buffer: HostBuffer[dtype]):
# Randomize
random.seed()
for i in range(len(buffer)):
buffer[i] = random.random_float64(
MATRIX_MIN_ELEM, MATRIX_MAX_ELEM
).cast[dtype]()[0]
def main():
try:
ctx = DeviceContext()
buffer_a = ctx.enqueue_create_buffer[dtype](
ROWS_A * COLS_A
).enqueue_fill(0.0)
buffer_b = ctx.enqueue_create_buffer[dtype](
ROWS_B * COLS_B
).enqueue_fill(0.0)
buffer_c = ctx.enqueue_create_buffer[dtype](
ROWS_C * COLS_C
).enqueue_fill(0.0)
with buffer_a.map_to_host() as h_buffer_a:
fill_buffer(h_buffer_a)
with buffer_b.map_to_host() as h_buffer_b:
fill_buffer(h_buffer_b)
matrix_a = MatrixA(buffer_a)
matrix_b = MatrixB(buffer_b)
matrix_c = MatrixC(buffer_c)
ctx.enqueue_function[
naive_matmul_one_thread_per_col[layout_a, layout_b, layout_c]
](
matrix_a,
matrix_b,
matrix_c,
grid_dim=BLOCKS,
block_dim=THREADS,
)
ctx.synchronize()
with buffer_a.map_to_host() as h_buffer_a:
with buffer_b.map_to_host() as h_buffer_b:
with buffer_c.map_to_host() as h_buffer_c:
assert_allclose(
(ROWS_A, COLS_A, h_buffer_a),
(ROWS_B, COLS_B, h_buffer_b),
(ROWS_C, COLS_C, h_buffer_c),
)
except e:
print("Prininting here: ", e)
def assert_allclose(
buff_a_with_dims: (Int, Int, HostBuffer[dtype]),
buff_b_with_dims: (Int, Int, HostBuffer[dtype]),
buff_c_with_dims: (Int, Int, HostBuffer[dtype]),
) raises:
a_rows, a_cols, a_buff = buff_a_with_dims
matrix_a = reshape(to_ndarray(a_buff), a_rows, a_cols)
b_rows, b_cols, b_buff = buff_b_with_dims
matrix_b = reshape(to_ndarray(b_buff), b_rows, b_cols)
c_rows, c_cols, c_buff = buff_c_with_dims
matrix_c = reshape(to_ndarray(c_buff), c_rows, c_cols)
np = Python.import_module("numpy")
assert_true(np.allclose(np.matmul(matrix_a, matrix_b), matrix_c))
print("Assertion was successful")
def to_ndarray(buffer: HostBuffer[dtype]) raises -> PythonObject:
np = Python.import_module("numpy")
ndarray = np.zeros(len(buffer), dtype=np.float32)
ndarray_ptr = ndarray_ptr[dtype](ndarray)
buffer_ptr = buffer.unsafe_ptr()
memcpy(ndarray_ptr, buffer_ptr, len(buffer))
return ndarray
def reshape(ndarray: PythonObject, rows: Int, cols: Int) raises -> PythonObject:
return ndarray.reshape(rows, cols)
def ndarray_ptr[
dtype: DType
](ndarray: PythonObject) raises -> UnsafePointer[Scalar[dtype]]:
return ndarray.__array_interface__["data"][0].unsafe_get_as_pointer[dtype]()
Dumb matrix multiplication
Simulate the CPU-style matrix multiplication with 1 GPU thread per row
from gpu.host import DeviceContext, HostBuffer
from gpu import thread_idx, block_idx, block_dim
import random
from layout import Layout, LayoutTensor
from memory import UnsafePointer, memcpy
from python import Python, PythonObject
from std.testing import assert_true
comptime ROWS_A = 64
comptime COLS_A = 16
comptime ROWS_B = 16
comptime COLS_B = 8
comptime ROWS_C = ROWS_A
comptime COLS_C = COLS_B
comptime MATRIX_MIN_ELEM = -5.0
comptime MATRIX_MAX_ELEM = 5.0
comptime dtype = DType.float32
# Num threads per block
comptime THREADS = ROWS_C
# Total numbers blocks in the grid
comptime BLOCKS = 1
comptime layout_a = Layout.row_major(ROWS_A, COLS_A)
comptime layout_b = Layout.row_major(ROWS_B, COLS_B)
comptime layout_c = Layout.row_major(ROWS_C, COLS_C)
comptime MatrixA = LayoutTensor[dtype, layout_a, MutableAnyOrigin]
comptime MatrixB = LayoutTensor[dtype, layout_b, MutableAnyOrigin]
comptime MatrixC = LayoutTensor[dtype, layout_c, MutableAnyOrigin]
def naive_matmul_one_thread_per_row[
a: Layout, b: Layout, c: Layout
](A: MatrixA, B: MatrixB, C: MatrixC,):
var tid = block_idx.x * block_dim.x + thread_idx.x
if tid < ROWS_A: # Each thread id `tid` is a row of A or C
for j in range(COLS_B):
for k in range(COLS_A):
C[tid, j] += A[tid, k] * B[k, j]
# Initialize the matrix buffer with values in the range 0 to 100
def fill_buffer(buffer: HostBuffer[dtype]):
# Randomize
random.seed()
for i in range(len(buffer)):
buffer[i] = random.random_float64(
MATRIX_MIN_ELEM, MATRIX_MAX_ELEM
).cast[dtype]()[0]
def main():
try:
ctx = DeviceContext()
buffer_a = ctx.enqueue_create_buffer[dtype](
ROWS_A * COLS_A
).enqueue_fill(0.0)
buffer_b = ctx.enqueue_create_buffer[dtype](
ROWS_B * COLS_B
).enqueue_fill(0.0)
buffer_c = ctx.enqueue_create_buffer[dtype](
ROWS_C * COLS_C
).enqueue_fill(0.0)
with buffer_a.map_to_host() as h_buffer_a:
fill_buffer(h_buffer_a)
with buffer_b.map_to_host() as h_buffer_b:
fill_buffer(h_buffer_b)
matrix_a = MatrixA(buffer_a)
matrix_b = MatrixB(buffer_b)
matrix_c = MatrixC(buffer_c)
ctx.enqueue_function[
naive_matmul_one_thread_per_row[layout_a, layout_b, layout_c]
](
matrix_a,
matrix_b,
matrix_c,
grid_dim=BLOCKS,
block_dim=THREADS,
)
ctx.synchronize()
with buffer_a.map_to_host() as h_buffer_a:
with buffer_b.map_to_host() as h_buffer_b:
with buffer_c.map_to_host() as h_buffer_c:
assert_allclose(
(ROWS_A, COLS_A, h_buffer_a),
(ROWS_B, COLS_B, h_buffer_b),
(ROWS_C, COLS_C, h_buffer_c),
)
except e:
print("Prininting here: ", e)
def assert_allclose(
buff_a_with_dims: (Int, Int, HostBuffer[dtype]),
buff_b_with_dims: (Int, Int, HostBuffer[dtype]),
buff_c_with_dims: (Int, Int, HostBuffer[dtype]),
) raises:
a_rows, a_cols, a_buff = buff_a_with_dims
matrix_a = reshape(to_ndarray(a_buff), a_rows, a_cols)
b_rows, b_cols, b_buff = buff_b_with_dims
matrix_b = reshape(to_ndarray(b_buff), b_rows, b_cols)
c_rows, c_cols, c_buff = buff_c_with_dims
matrix_c = reshape(to_ndarray(c_buff), c_rows, c_cols)
np = Python.import_module("numpy")
assert_true(np.allclose(np.matmul(matrix_a, matrix_b), matrix_c))
print("Assertion was successful")
def to_ndarray(buffer: HostBuffer[dtype]) raises -> PythonObject:
np = Python.import_module("numpy")
ndarray = np.zeros(len(buffer), dtype=np.float32)
ndarray_ptr = ndarray_ptr[dtype](ndarray)
buffer_ptr = buffer.unsafe_ptr()
memcpy(ndarray_ptr, buffer_ptr, len(buffer))
return ndarray
def reshape(ndarray: PythonObject, rows: Int, cols: Int) raises -> PythonObject:
return ndarray.reshape(rows, cols)
def ndarray_ptr[
dtype: DType
](ndarray: PythonObject) raises -> UnsafePointer[Scalar[dtype]]:
return ndarray.__array_interface__["data"][0].unsafe_get_as_pointer[dtype]()
Dumb matrix multiplication
Simulate the CPU-style triple for-loop truly dumb matrix multiplication
from gpu.host import DeviceContext, HostBuffer
from gpu import thread_idx, block_idx, block_dim
import random
from layout import Layout, LayoutTensor
from memory import UnsafePointer, memcpy
from python import Python, PythonObject
from std.testing import assert_true
comptime ROWS_A = 64
comptime COLS_A = 16
comptime ROWS_B = 16
comptime COLS_B = 8
comptime ROWS_C = ROWS_A
comptime COLS_C = COLS_B
comptime MATRIX_MIN_ELEM = -5.0
comptime MATRIX_MAX_ELEM = 5.0
comptime dtype = DType.float32
# Num threads per block
comptime THREADS = 1
# Total numbers blocks in the grid
comptime BLOCKS = 1
comptime layout_a = Layout.row_major(ROWS_A, COLS_A)
comptime layout_b = Layout.row_major(ROWS_B, COLS_B)
comptime layout_c = Layout.row_major(ROWS_C, COLS_C)
comptime MatrixA = LayoutTensor[dtype, layout_a, MutableAnyOrigin]
comptime MatrixB = LayoutTensor[dtype, layout_b, MutableAnyOrigin]
comptime MatrixC = LayoutTensor[dtype, layout_c, MutableAnyOrigin]
def naive_matmul_single_thread_layout_tensor[
a: Layout, b: Layout, c: Layout
](A: MatrixA, B: MatrixB, C: MatrixC,):
var tid = block_idx.x * block_dim.x + thread_idx.x
if tid == 0:
for i in range(ROWS_A):
for j in range(COLS_B):
for k in range(COLS_A):
C[i, j] += A[i, k] * B[k, j]
# Initialize the matrix buffer with values in the range 0 to 100
def fill_buffer(buffer: HostBuffer[dtype]):
# Randomize
random.seed()
for i in range(len(buffer)):
buffer[i] = random.random_float64(
MATRIX_MIN_ELEM, MATRIX_MAX_ELEM
).cast[dtype]()[0]
def main():
try:
ctx = DeviceContext()
buffer_a = ctx.enqueue_create_buffer[dtype](
ROWS_A * COLS_A
).enqueue_fill(0.0)
buffer_b = ctx.enqueue_create_buffer[dtype](
ROWS_B * COLS_B
).enqueue_fill(0.0)
buffer_c = ctx.enqueue_create_buffer[dtype](
ROWS_C * COLS_C
).enqueue_fill(0.0)
with buffer_a.map_to_host() as h_buffer_a:
fill_buffer(h_buffer_a)
with buffer_b.map_to_host() as h_buffer_b:
fill_buffer(h_buffer_b)
matrix_a = MatrixA(buffer_a)
matrix_b = MatrixB(buffer_b)
matrix_c = MatrixC(buffer_c)
ctx.enqueue_function[
naive_matmul_single_thread_layout_tensor[
layout_a, layout_b, layout_c
]
](
matrix_a,
matrix_b,
matrix_c,
grid_dim=BLOCKS,
block_dim=THREADS,
)
ctx.synchronize()
with buffer_a.map_to_host() as h_buffer_a:
with buffer_b.map_to_host() as h_buffer_b:
with buffer_c.map_to_host() as h_buffer_c:
assert_allclose(
(ROWS_A, COLS_A, h_buffer_a),
(ROWS_B, COLS_B, h_buffer_b),
(ROWS_C, COLS_C, h_buffer_c),
)
except e:
print("Prininting here: ", e)
def assert_allclose(
buff_a_with_dims: (Int, Int, HostBuffer[dtype]),
buff_b_with_dims: (Int, Int, HostBuffer[dtype]),
buff_c_with_dims: (Int, Int, HostBuffer[dtype]),
) raises:
a_rows, a_cols, a_buff = buff_a_with_dims
matrix_a = reshape(to_ndarray(a_buff), a_rows, a_cols)
b_rows, b_cols, b_buff = buff_b_with_dims
matrix_b = reshape(to_ndarray(b_buff), b_rows, b_cols)
c_rows, c_cols, c_buff = buff_c_with_dims
matrix_c = reshape(to_ndarray(c_buff), c_rows, c_cols)
np = Python.import_module("numpy")
assert_true(np.allclose(np.matmul(matrix_a, matrix_b), matrix_c))
print("Assertion was successful")
def to_ndarray(buffer: HostBuffer[dtype]) raises -> PythonObject:
np = Python.import_module("numpy")
ndarray = np.zeros(len(buffer), dtype=np.float32)
ndarray_ptr = ndarray_ptr[dtype](ndarray)
buffer_ptr = buffer.unsafe_ptr()
memcpy(ndarray_ptr, buffer_ptr, len(buffer))
return ndarray
def reshape(ndarray: PythonObject, rows: Int, cols: Int) raises -> PythonObject:
return ndarray.reshape(rows, cols)
def ndarray_ptr[
dtype: DType
](ndarray: PythonObject) raises -> UnsafePointer[Scalar[dtype]]:
return ndarray.__array_interface__["data"][0].unsafe_get_as_pointer[dtype]()
Dumb matrix multiplication
Use one one GPU thread for each column of the output matrix
from gpu.host import DeviceContext, HostBuffer
from gpu import thread_idx, block_idx, block_dim
import random
from layout import Layout, LayoutTensor
from memory import UnsafePointer, memcpy
from python import Python, PythonObject
from std.testing import assert_true
comptime ROWS_A = 64
comptime COLS_A = 16
comptime ROWS_B = 16
comptime COLS_B = 8
comptime ROWS_C = ROWS_A
comptime COLS_C = COLS_B
comptime MATRIX_MIN_ELEM = -5.0
comptime MATRIX_MAX_ELEM = 5.0
comptime dtype = DType.float32
# Num threads per block
comptime THREADS = (5, 5)
# Total numbers blocks in the grid
comptime BLOCKS = (
(COLS_C + THREADS[0] - 1) // THREADS[0],
(ROWS_C + THREADS[1] - 1) // THREADS[1],
)
comptime layout_a = Layout.row_major(ROWS_A, COLS_A)
comptime layout_b = Layout.row_major(ROWS_B, COLS_B)
comptime layout_c = Layout.row_major(ROWS_C, COLS_C)
comptime MatrixA = LayoutTensor[dtype, layout_a, MutableAnyOrigin]
comptime MatrixB = LayoutTensor[dtype, layout_b, MutableAnyOrigin]
comptime MatrixC = LayoutTensor[dtype, layout_c, MutableAnyOrigin]
def matmul_thread_per_output_cell[
a: Layout, b: Layout, c: Layout
](A: MatrixA, B: MatrixB, C: MatrixC,):
var i = block_idx.y * block_dim.y + thread_idx.y # Rows
var j = block_idx.x * block_dim.x + thread_idx.x # Colums
if i < ROWS_C and j < COLS_C:
for k in range(ROWS_B):
C[i, j] += A[i, k] * B[k, j]
# Initialize the matrix buffer with values in the range 0 to 100
def fill_buffer(buffer: HostBuffer[dtype]):
# Randomize
random.seed()
for i in range(len(buffer)):
buffer[i] = random.random_float64(
MATRIX_MIN_ELEM, MATRIX_MAX_ELEM
).cast[dtype]()[0]
def main():
try:
ctx = DeviceContext()
buffer_a = ctx.enqueue_create_buffer[dtype](
ROWS_A * COLS_A
).enqueue_fill(0.0)
buffer_b = ctx.enqueue_create_buffer[dtype](
ROWS_B * COLS_B
).enqueue_fill(0.0)
buffer_c = ctx.enqueue_create_buffer[dtype](
ROWS_C * COLS_C
).enqueue_fill(0.0)
with buffer_a.map_to_host() as h_buffer_a:
fill_buffer(h_buffer_a)
with buffer_b.map_to_host() as h_buffer_b:
fill_buffer(h_buffer_b)
matrix_a = MatrixA(buffer_a)
matrix_b = MatrixB(buffer_b)
matrix_c = MatrixC(buffer_c)
ctx.enqueue_function[
matmul_thread_per_output_cell[layout_a, layout_b, layout_c]
](
matrix_a,
matrix_b,
matrix_c,
grid_dim=BLOCKS,
block_dim=THREADS,
)
ctx.synchronize()
with buffer_a.map_to_host() as h_buffer_a:
with buffer_b.map_to_host() as h_buffer_b:
with buffer_c.map_to_host() as h_buffer_c:
assert_allclose(
(ROWS_A, COLS_A, h_buffer_a),
(ROWS_B, COLS_B, h_buffer_b),
(ROWS_C, COLS_C, h_buffer_c),
)
except e:
print("Prininting here: ", e)
def assert_allclose(
buff_a_with_dims: (Int, Int, HostBuffer[dtype]),
buff_b_with_dims: (Int, Int, HostBuffer[dtype]),
buff_c_with_dims: (Int, Int, HostBuffer[dtype]),
) raises:
a_rows, a_cols, a_buff = buff_a_with_dims
matrix_a = reshape(to_ndarray(a_buff), a_rows, a_cols)
b_rows, b_cols, b_buff = buff_b_with_dims
matrix_b = reshape(to_ndarray(b_buff), b_rows, b_cols)
c_rows, c_cols, c_buff = buff_c_with_dims
matrix_c = reshape(to_ndarray(c_buff), c_rows, c_cols)
np = Python.import_module("numpy")
assert_true(np.allclose(np.matmul(matrix_a, matrix_b), matrix_c))
print("Assertion was successful")
def to_ndarray(buffer: HostBuffer[dtype]) raises -> PythonObject:
np = Python.import_module("numpy")
ndarray = np.zeros(len(buffer), dtype=np.float32)
ndarray_ptr = ndarray_ptr[dtype](ndarray)
buffer_ptr = buffer.unsafe_ptr()
memcpy(ndarray_ptr, buffer_ptr, len(buffer))
return ndarray
def reshape(ndarray: PythonObject, rows: Int, cols: Int) raises -> PythonObject:
return ndarray.reshape(rows, cols)
def ndarray_ptr[
dtype: DType
](ndarray: PythonObject) raises -> UnsafePointer[Scalar[dtype]]:
return ndarray.__array_interface__["data"][0].unsafe_get_as_pointer[dtype]()
Use one one GPU thread for each column of the output matrix
Uses shared memory via stack_allocation
from gpu.host import DeviceContext, HostBuffer
from gpu import thread_idx, block_idx, block_dim
import random
from layout import Layout, LayoutTensor
from memory import UnsafePointer, memcpy, stack_allocation
from python import Python, PythonObject
from std.testing import assert_true
from algorithm import vectorize
from sys import simdwidthof, strided_load
comptime ROWS_A = 9
comptime COLS_A = 17
comptime ROWS_B = 17
comptime COLS_B = 7
comptime ROWS_C = ROWS_A
comptime COLS_C = COLS_B
comptime MATRIX_MIN_ELEM = -5.0
comptime MATRIX_MAX_ELEM = 5.0
comptime dtype = DType.float32
# Num threads per block
comptime THREADS = (5, 5)
# Total numbers blocks in the grid
comptime BLOCKS = (
(COLS_C + THREADS[0] - 1) // THREADS[0],
(ROWS_C + THREADS[1] - 1) // THREADS[1],
)
comptime layout_a = Layout.row_major(ROWS_A, COLS_A)
comptime layout_b = Layout.row_major(ROWS_B, COLS_B)
comptime layout_c = Layout.row_major(ROWS_C, COLS_C)
comptime MatrixA = LayoutTensor[dtype, layout_a, MutableAnyOrigin]
comptime MatrixB = LayoutTensor[dtype, layout_b, MutableAnyOrigin]
comptime MatrixC = LayoutTensor[dtype, layout_c, MutableAnyOrigin]
comptime Storage = LayoutTensor[
dtype, Layout.row_major(1, simdwidthof[dtype]()), MutableAnyOrigin
]
def matmul_thread_per_output_cell_vectorized(
A: MatrixA, B: MatrixB, C: MatrixC, store: Storage
):
var i = block_idx.y * block_dim.y + thread_idx.y # Rows
var j = block_idx.x * block_dim.x + thread_idx.x # Colums
if i < ROWS_C and j < COLS_C:
tile = stack_allocation[ROWS_B, Scalar[dtype]]()
each_b_col = B.tile[ROWS_B, 1](0, j)
for k in range(ROWS_B):
tile[k] = each_b_col[k, 0][0]
@parameter
def dotproduct[simd_width: Int](idx: Int):
C[i, j] += (
A.load[width=simd_width](i, idx)
* tile.load[width=simd_width](idx)
).reduce_add()
vectorize[dotproduct, simdwidthof[dtype]()](ROWS_B)
# Initialize the matrix buffer with values in the range 0 to 100
def fill_buffer(buffer: HostBuffer[dtype]):
# Randomize
random.seed()
for i in range(len(buffer)):
buffer[i] = random.random_float64(
MATRIX_MIN_ELEM, MATRIX_MAX_ELEM
).cast[dtype]()[0]
def main():
try:
ctx = DeviceContext()
buffer_a = ctx.enqueue_create_buffer[dtype](
ROWS_A * COLS_A
).enqueue_fill(0.0)
buffer_b = ctx.enqueue_create_buffer[dtype](
ROWS_B * COLS_B
).enqueue_fill(0.0)
buffer_c = ctx.enqueue_create_buffer[dtype](
ROWS_C * COLS_C
).enqueue_fill(0.0)
store = ctx.enqueue_create_buffer[dtype](
simdwidthof[dtype]()
).enqueue_fill(0.0)
with buffer_a.map_to_host() as h_buffer_a:
fill_buffer(h_buffer_a)
with buffer_b.map_to_host() as h_buffer_b:
fill_buffer(h_buffer_b)
matrix_a = MatrixA(buffer_a)
matrix_b = MatrixB(buffer_b)
matrix_c = MatrixC(buffer_c)
storage = Storage(store)
ctx.enqueue_function[matmul_thread_per_output_cell_vectorized](
matrix_a,
matrix_b,
matrix_c,
storage,
grid_dim=BLOCKS,
block_dim=THREADS,
)
ctx.synchronize()
with buffer_a.map_to_host() as h_buffer_a:
with buffer_b.map_to_host() as h_buffer_b:
with buffer_c.map_to_host() as h_buffer_c:
assert_allclose(
(ROWS_A, COLS_A, h_buffer_a),
(ROWS_B, COLS_B, h_buffer_b),
(ROWS_C, COLS_C, h_buffer_c),
)
except e:
print("Prininting here: ", e)
def assert_allclose(
buff_a_with_dims: (Int, Int, HostBuffer[dtype]),
buff_b_with_dims: (Int, Int, HostBuffer[dtype]),
buff_c_with_dims: (Int, Int, HostBuffer[dtype]),
) raises:
a_rows, a_cols, a_buff = buff_a_with_dims
matrix_a = reshape(to_ndarray(a_buff), a_rows, a_cols)
b_rows, b_cols, b_buff = buff_b_with_dims
matrix_b = reshape(to_ndarray(b_buff), b_rows, b_cols)
c_rows, c_cols, c_buff = buff_c_with_dims
matrix_c = reshape(to_ndarray(c_buff), c_rows, c_cols)
np = Python.import_module("numpy")
assert_true(np.allclose(np.matmul(matrix_a, matrix_b), matrix_c))
print("Assertion was successful")
def to_ndarray(buffer: HostBuffer[dtype]) raises -> PythonObject:
np = Python.import_module("numpy")
ndarray = np.zeros(len(buffer), dtype=np.float32)
ndarray_ptr = ndarray_ptr[dtype](ndarray)
buffer_ptr = buffer.unsafe_ptr()
memcpy(ndarray_ptr, buffer_ptr, len(buffer))
return ndarray
def reshape(ndarray: PythonObject, rows: Int, cols: Int) raises -> PythonObject:
return ndarray.reshape(rows, cols)
def ndarray_ptr[
dtype: DType
](ndarray: PythonObject) raises -> UnsafePointer[Scalar[dtype]]:
return ndarray.__array_interface__["data"][0].unsafe_get_as_pointer[dtype]()
Utils
### Timer utility
### A simple RAII / context-manager timer that prints elapsed wall-clock time in nanoseconds.
from std.time import global_perf_counter_ns
# Timer struct used via Python-style `with` blocks.
#
# Records the time at `__enter__` and prints the elapsed duration at `__exit__`.
# The label prefix is a templated `StringSlice` to avoid allocation.
#
# Usage:
# with Timer("My computation: "):
# do_something()
# # Prints: "My computation: 12345000 nanoseconds"
#
@fieldwise_init
struct Timer[origin: Origin, //](ImplicitlyCopyable):
var start_time: UInt64
var prefix: StringSlice[Self.origin]
def __init__(out self, prefix: StringSlice[Self.origin]):
self.start_time = 0
self.prefix = prefix
def __enter__(mut self) -> Self:
self.start_time = global_perf_counter_ns()
return self
def __exit__(mut self):
elapsed_time_ms = global_perf_counter_ns() - self.start_time
print(self.prefix, elapsed_time_ms, "nanoseconds")
Histogram
Program to compute histogram of a 1D array
from gpu.host import DeviceContext, HostBuffer, DeviceBuffer
from gpu import thread_idx, block_idx, block_dim
import random
from math import ceildiv
from memory import UnsafePointer
from layout import Layout, LayoutTensor
from os import Atomic
from os.atomic import Consistency
comptime dtype = DType.int64
# How many numbers to bin? 2 ^ 20 (default)
comptime ELEMS_COUNT = 1 << 20
# How many bins?
comptime NUM_BINS = 10
# Num threads per block
comptime THREADS = 256
# Total numbers blocks in the grid
comptime BLOCKS = ceildiv(ELEMS_COUNT, THREADS)
# Max value of any binned element
comptime MAX_ELEM = 101
comptime MIN_ELEM = 1
comptime BIN_WIDTH = (MAX_ELEM - MIN_ELEM + 1) // NUM_BINS
comptime input_layout = Layout.row_major(ELEMS_COUNT)
def histogram(
input: LayoutTensor[dtype, input_layout, MutableAnyOrigin],
output: UnsafePointer[Scalar[dtype]],
total_elems: Int,
):
var tid = block_idx.x * block_dim.x + thread_idx.x
if tid < total_elems:
var elem = input[tid]
bin_index = bin_index(elem[0])
# _ = Atomic.fetch_add[ordering= Consistency.MONOTONIC](output + bin_index, 1)
_ = Atomic.fetch_add(output + bin_index, 1)
# Initialize the input buffer with values in the range 0 to 100
def fill_buffer(buffer: HostBuffer[dtype]):
# Randomize
random.seed()
for i in range(len(buffer)):
buffer[i] = random.random_ui64(MIN_ELEM, MAX_ELEM).cast[dtype]()[0]
# Find the bin index given a number
@always_inline
def bin_index(elem: Int64) -> Int:
bin_index = Int((elem - MIN_ELEM) // BIN_WIDTH)
if bin_index >= NUM_BINS:
bin_index = NUM_BINS - 1
elif bin_index < 0:
bin_index = 0
return bin_index
def main():
try:
ctx = DeviceContext()
elements = ctx.enqueue_create_buffer[dtype](ELEMS_COUNT)
bins = ctx.enqueue_create_buffer[dtype](NUM_BINS).enqueue_fill(0)
with elements.map_to_host() as host_elements:
fill_buffer(host_elements)
input_tensor = LayoutTensor[dtype, input_layout, MutableAnyOrigin](
elements
)
# output_tensor = LayoutTensor[mut=True, dtype, output_layout](bins)
ctx.enqueue_function[histogram](
input_tensor,
bins.unsafe_ptr(),
ELEMS_COUNT,
grid_dim=BLOCKS,
block_dim=THREADS,
)
ctx.synchronize()
with bins.map_to_host() as bins_host:
print(bins_host)
print(ctx.name())
except e:
print("Prininting here: ", e)