reasoning-gym/reasoning_gym/arc/arc_1d_tasks.py

1227 lines
40 KiB
Python

from random import Random
from typing import Optional
def gen_field(size: int, color: int = 0) -> list[int]:
"""Generate a field of given size filled with specified color (default 0)."""
return [color] * size
def write_block(pos: int, block: list[int], field: list[int]) -> list[int]:
"""Write a block into a field at given position."""
result = field.copy()
for i, color in enumerate(block):
result[pos + i] = color
return result
def task_move_n_pix(rng: Random, size: int, move_pix: int, solid: bool) -> Optional[dict[str, list[int]]]:
"""Generate a task where a block is moved to the right by move_pix pixels."""
if size <= move_pix + 1:
return None
block_size = rng.randint(1, size - move_pix - 1)
block_pos = rng.randint(0, size - block_size - move_pix)
if solid:
color = rng.randint(1, 9)
block = [color] * block_size
else:
block = [rng.randint(1, 9) for _ in range(block_size)]
question = write_block(block_pos, block, gen_field(size))
answer = write_block(block_pos + move_pix, block, gen_field(size))
return {"input": question, "output": answer}
def task_move_n_pix_wrapped(rng: Random, size: int, move_pix: int, solid: bool) -> Optional[dict[str, list[int]]]:
"""Generate a task where a block is moved to the right by move_pix pixels with wrapping."""
block_size = rng.randint(1, size)
block_pos = rng.randint(0, size)
if solid:
color = rng.randint(1, 9)
block = [color] * block_size
else:
block = [rng.randint(1, 9) for _ in range(block_size)]
question = gen_field(size)
answer = gen_field(size)
for i, color in enumerate(block):
question[(block_pos + i) % size] = color
answer[(block_pos + move_pix + i) % size] = color
return {"input": question, "output": answer}
def task_gravity(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where all non-zero elements are attracted to the left."""
density = 0.5
question = [rng.randint(1, 9) if rng.random() < density else 0 for _ in range(size)]
non_zero = [x for x in question if x != 0]
answer = non_zero + [0] * (size - len(non_zero))
return {"input": question, "output": answer}
def task_gravity_counting(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where non-zero elements are counted and represented as a sequence of 1s."""
density = 0.5
question = [rng.randint(1, 9) if rng.random() < density else 0 for _ in range(size)]
count = sum(1 for x in question if x != 0)
answer = [1] * count + [0] * (size - count)
return {"input": question, "output": answer}
def task_gravity_antigravity(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where color 1 moves right and color 2 moves left."""
density = 0.5
question = [rng.randint(1, 2) if rng.random() < density else 0 for _ in range(size)]
color1 = [x for x in question if x == 1]
color2 = [x for x in question if x == 2]
answer = [2] * len(color2) + [0] * (size - len(color1) - len(color2)) + [1] * len(color1)
return {"input": question, "output": answer}
def task_block_touch_dot(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a block moves to touch (but not cover) a dot."""
dot_color = 1
block_color = rng.randint(2, 9)
block_size = rng.randint(1, size)
dot_pos = rng.randint(0, size)
can_place_left = dot_pos >= block_size
can_place_right = dot_pos + block_size < size
if not (can_place_left or can_place_right):
return None
if can_place_left and can_place_right:
side = rng.choice(["left", "right"])
elif can_place_left:
side = "left"
else:
side = "right"
if side == "left":
q_block_pos = rng.randint(0, dot_pos - block_size)
a_block_pos = dot_pos - block_size
else:
q_block_pos = rng.randint(dot_pos + 1, size - block_size)
a_block_pos = dot_pos + 1
question = gen_field(size)
question[dot_pos] = dot_color
question = write_block(q_block_pos, [block_color] * block_size, question)
answer = gen_field(size)
answer[dot_pos] = dot_color
answer = write_block(a_block_pos, [block_color] * block_size, answer)
return {"input": question, "output": answer}
def task_block_touch_dot_n_pix(rng: Random, size: int, move_pix: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a block moves move_pix pixels toward a dot."""
dot_color = 2
block_color = rng.randint(3, 9)
block_size = rng.randint(1, size)
dot_pos = rng.randint(0, size)
can_place_left = dot_pos >= block_size
can_place_right = dot_pos + block_size < size
if not (can_place_left or can_place_right):
return None
if can_place_left and can_place_right:
side = rng.choice(["left", "right"])
elif can_place_left:
side = "left"
else:
side = "right"
if side == "left":
q_block_pos = rng.randint(0, dot_pos - block_size)
distance = (dot_pos - block_size) - q_block_pos
move = min(distance, move_pix)
a_block_pos = q_block_pos + move
else:
q_block_pos = rng.randint(dot_pos + 1, size - block_size)
distance = q_block_pos - (dot_pos + 1)
move = min(distance, move_pix)
a_block_pos = q_block_pos - move
question = gen_field(size)
question[dot_pos] = dot_color
question = write_block(q_block_pos, [block_color] * block_size, question)
answer = gen_field(size)
answer[dot_pos] = dot_color
answer = write_block(a_block_pos, [block_color] * block_size, answer)
return {"input": question, "output": answer}
def task_block_scale_to_dot(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a block scales to touch a dot (keeping one end fixed)."""
dot_color = 2
block_color = rng.randint(3, 9)
block_size = rng.randint(1, size)
dot_pos = rng.randint(0, size)
can_place_left = dot_pos >= block_size
can_place_right = dot_pos + block_size < size
if not (can_place_left or can_place_right):
return None
if can_place_left and can_place_right:
side = rng.choice(["left", "right"])
elif can_place_left:
side = "left"
else:
side = "right"
if side == "left":
q_block_pos = rng.randint(0, dot_pos - block_size)
new_size = dot_pos - q_block_pos + 1
a_block_pos = q_block_pos
else:
q_block_pos = rng.randint(dot_pos + 1, size - block_size)
new_size = (q_block_pos + block_size) - dot_pos
a_block_pos = dot_pos
question = gen_field(size)
question[dot_pos] = dot_color
question = write_block(q_block_pos, [block_color] * block_size, question)
answer = gen_field(size)
answer[dot_pos] = dot_color
answer = write_block(a_block_pos, [block_color] * new_size, answer)
return {"input": question, "output": answer}
def task_two_points_and_fill(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where space between two points of same color is filled with that color."""
color = rng.randint(1, 9)
pos1 = rng.randint(0, size - 1)
pos2 = rng.randint(0, size - 1)
if pos1 == pos2:
return None
pos1, pos2 = min(pos1, pos2), max(pos1, pos2)
question = gen_field(size)
question[pos1] = color
question[pos2] = color
answer = question.copy()
for i in range(pos1, pos2 + 1):
answer[i] = color
return {"input": question, "output": answer}
def task_reflect_block_with_border_pixel(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a block with a border pixel is reflected."""
block_size = rng.randint(2, size)
if block_size > size:
return None
c1 = rng.randint(1, 9)
c2 = rng.randint(1, 9)
if c1 == c2:
return None
side = "left" if rng.random() < 0.5 else "right"
pos = rng.randint(0, size - block_size)
block = [c1] * block_size
if side == "left":
block[0] = c2
else:
block[block_size - 1] = c2
question = write_block(pos, block, gen_field(size))
reversed_block = block[::-1] # Reverse the block
answer = write_block(pos, reversed_block, gen_field(size))
return {"input": question, "output": answer}
def task_reflect_block_with_border_pixel_random(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a random-colored block with a border pixel is reflected."""
block_size = rng.randint(2, size)
if block_size > size:
return None
side = "left" if rng.random() < 0.5 else "right"
pos = rng.randint(0, size - block_size)
block = [rng.randint(1, 9) for _ in range(block_size)]
border_color = rng.randint(1, 9)
if side == "left":
if block[0] == border_color:
return None
block[0] = border_color
else:
if block[block_size - 1] == border_color:
return None
block[block_size - 1] = border_color
question = write_block(pos, block, gen_field(size))
reversed_block = block[::-1] # Reverse the block
answer = write_block(pos, reversed_block, gen_field(size))
return {"input": question, "output": answer}
def task_reflect_block_around_dot(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a block is reflected around a dot."""
dot_color = 2
dot_pos = rng.randint(0, size)
block_size = rng.randint(1, size)
block_pos = rng.randint(0, size - block_size)
block_end = block_pos + block_size - 1
# Check if block is strictly to left or right of dot
strictly_left = block_end < dot_pos
strictly_right = block_pos > dot_pos
if not (strictly_left or strictly_right):
return None
block_color = rng.randint(3, 9) # Different from dot color
block = [block_color] * block_size
# Calculate reflection bounds
min_reflect = 2 * dot_pos - block_end
max_reflect = 2 * dot_pos - block_pos
if min_reflect < 0 or max_reflect >= size:
return None
question = gen_field(size)
question = write_block(block_pos, block, question)
question[dot_pos] = dot_color
answer = gen_field(size)
answer[dot_pos] = dot_color
for i in range(block_size):
reflect_idx = 2 * dot_pos - (block_pos + i)
answer[reflect_idx] = block[i]
return {"input": question, "output": answer}
def task_block_and_noise_remove(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where noise around a block needs to be removed."""
block_size = rng.randint(2, size)
if block_size > size:
return None
block_pos = rng.randint(0, size - block_size)
color = rng.randint(1, 9)
# Create field with block
field = gen_field(size)
for i in range(block_size):
field[block_pos + i] = color
# Track forbidden positions for noise
forbidden = [False] * size
for i in range(block_pos, block_pos + block_size):
forbidden[i] = True
if block_pos > 0:
forbidden[block_pos - 1] = True
if block_pos + block_size < size:
forbidden[block_pos + block_size] = True
# Add noise
noise_count = rng.randint(1, 3)
noise_positions = []
for _ in range(noise_count):
allowed = [i for i in range(size) if not forbidden[i]]
if not allowed:
break
noise_pos = rng.choice(allowed)
noise_positions.append(noise_pos)
field[noise_pos] = color
forbidden[noise_pos] = True
if noise_pos > 0:
forbidden[noise_pos - 1] = True
if noise_pos + 1 < size:
forbidden[noise_pos + 1] = True
if len(noise_positions) < noise_count:
return None
question = field
answer = field.copy()
for pos in noise_positions:
answer[pos] = 0
return {"input": question, "output": answer}
def task_block_and_noise_remove_inside(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where noise inside a block needs to be removed."""
if size <= 6:
return None
block_size = rng.randint(6, size)
if block_size > size:
return None
block_pos = rng.randint(0, size - block_size)
color = rng.randint(1, 9)
# Create field with block
field = gen_field(size)
for i in range(block_size):
field[block_pos + i] = color
# Add noise inside block
max_noise = max(1, (block_size // 2) - 1)
noise_count = rng.randint(1, max_noise)
positions = list(range(block_size))
rng.shuffle(positions)
noise_positions = positions[:noise_count]
for offset in noise_positions:
pos = block_pos + offset
noise_color = rng.randint(1, 9)
while noise_color == color:
noise_color = rng.randint(1, 9)
field[pos] = noise_color
question = field
answer = field.copy()
for offset in noise_positions:
answer[block_pos + offset] = color
return {"input": question, "output": answer}
def task_copy_block_to_dots(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a block pattern is copied to dot positions."""
block_size = 3 if rng.random() < 0.5 else 5
if block_size >= size:
return None
color = rng.randint(1, 9)
block = [color] * block_size
# Generate dots with minimum distance to prevent overlap
min_gap = block_size
dot_positions = []
pos = block_size + block_size // 2 + 1
while pos <= size - block_size:
if rng.random() < 0.5: # Control dot density
dot_positions.append(pos)
pos += min_gap
pos += 1
if not dot_positions:
return None
question = gen_field(size)
question = write_block(0, block, question)
for pos in dot_positions:
question[pos] = color
answer = gen_field(size)
answer = write_block(0, block, answer)
for pos in dot_positions:
block_start = pos - block_size // 2
answer = write_block(block_start, block, answer)
return {"input": question, "output": answer}
def task_copy_block_to_dots_colors(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a block pattern is copied to dot positions with matching colors."""
block_size = 3 if rng.random() < 0.5 else 5
if block_size >= size:
return None
block_color = rng.randint(1, 9)
block = [block_color] * block_size
# Generate dots with minimum distance to prevent overlap
min_gap = block_size
dot_positions = []
dot_colors = []
pos = block_size + block_size // 2 + 1
while pos < size - block_size:
if rng.random() < 0.5:
dot_color = rng.randint(1, 9)
dot_positions.append(pos)
dot_colors.append(dot_color)
pos += min_gap
pos += 1
if not dot_positions:
return None
question = gen_field(size)
question = write_block(0, block, question)
for i, pos in enumerate(dot_positions):
question[pos] = dot_colors[i]
answer = gen_field(size)
answer = write_block(0, block, answer)
for i, pos in enumerate(dot_positions):
block_start = pos - block_size // 2
colored_block = [dot_colors[i]] * block_size
answer = write_block(block_start, colored_block, answer)
return {"input": question, "output": answer}
def task_paint_biggest_block(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where the largest block is painted a different color."""
target_color = 1
initial_color = rng.randint(2, 9)
# Generate random blocks
question = gen_field(size)
blocks = []
pos = 0
while pos < size:
if rng.random() < 0.4 and size - pos >= 2:
block_size = rng.randint(2, min(size - pos, 6))
blocks.append((pos, block_size))
for i in range(block_size):
question[pos + i] = initial_color
pos += block_size + 1
else:
pos += 1
if len(blocks) < 2:
return None
# Find biggest block
biggest_pos, biggest_size = max(blocks, key=lambda x: x[1])
# Check if there are multiple blocks of the same size
biggest_count = sum(1 for _, size in blocks if size == biggest_size)
if biggest_count > 1:
return None
answer = question.copy()
for i in range(biggest_size):
answer[biggest_pos + i] = target_color
return {"input": question, "output": answer}
def task_sort_blocks_by_size(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where blocks are sorted by size with 1 pixel gaps."""
color = rng.randint(1, 9)
blocks = []
pos = 0
# Generate random blocks with random sizes
while pos < size:
if rng.random() < 0.4 and size - pos >= 2:
block_size = rng.randint(1, min(size - pos, 6))
blocks.append((pos, block_size))
pos += block_size + rng.randint(1, 4) # Random gaps
else:
pos += 1
if len(blocks) < 2:
return None
# Create input field
question = gen_field(size)
for pos, block_size in blocks:
for i in range(block_size):
question[pos + i] = color
# Sort blocks by size
blocks.sort(key=lambda x: x[1])
# Check if sorted blocks fit with gaps
total_space = sum(size for _, size in blocks) + len(blocks) - 1
if total_space > size:
return None
# Create answer field with sorted blocks
answer = gen_field(size)
current_pos = 0
for _, block_size in blocks:
for i in range(block_size):
answer[current_pos + i] = color
current_pos += block_size + 1 # One pixel gap
return {"input": question, "output": answer}
def task_sort_complete_sequence(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a complete sequence of block sizes is sorted."""
# Calculate max possible block size given total array size
max_size = 1
total_space = 0
while total_space + max_size + 1 <= size:
total_space += max_size + 1
max_size += 1
max_size -= 1
if max_size < 2:
return None
color = rng.randint(1, 9)
# Create sequence of all sizes from 1 to max_size
blocks = list(range(1, max_size + 1))
rng.shuffle(blocks)
# Create input field with shuffled blocks
question = gen_field(size)
pos = 0
for block_size in blocks:
for i in range(block_size):
question[pos + i] = color
pos += block_size + 1
# Create answer field with sorted blocks
answer = gen_field(size)
pos = 0
for block_size in range(1, max_size + 1):
for i in range(block_size):
answer[pos + i] = color
pos += block_size + 1
return {"input": question, "output": answer}
def task_recolor_blocks_by_size(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where two blocks are recolored based on their size."""
# Generate two different random sizes
size1 = rng.randint(2, 8)
size2 = rng.randint(2, 8)
while size2 == size1:
size2 = rng.randint(2, 8)
# Ensure both blocks fit with at least 1 gap
if size1 + size2 + 1 > size:
return None
# Place blocks with gap
pos1 = rng.randint(0, size - (size1 + size2 + 1))
pos2 = rng.randint(pos1 + size1 + 1, size - size2)
# Create input field with both blocks color 3
question = gen_field(size)
for i in range(size1):
question[pos1 + i] = 3
for i in range(size2):
question[pos2 + i] = 3
# Create answer field with recolored blocks
answer = question.copy()
if size1 > size2:
for i in range(size1):
answer[pos1 + i] = 1
for i in range(size2):
answer[pos2 + i] = 2
else:
for i in range(size1):
answer[pos1 + i] = 2
for i in range(size2):
answer[pos2 + i] = 1
return {"input": question, "output": answer}
def task_gravity_one_step(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where non-zero elements move one step left if possible."""
question = [rng.randint(1, 9) if rng.random() < 0.5 else 0 for _ in range(size)]
answer = question.copy()
# Move each non-zero pixel one step left if possible
for i in range(1, size):
if answer[i] != 0 and answer[i - 1] == 0:
answer[i - 1] = answer[i]
answer[i] = 0
return {"input": question, "output": answer}
def task_move_block_by_own_size(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a block moves right by its own size."""
block_size = rng.randint(1, size // 2) # Ensure space for movement
pos = rng.randint(0, size - block_size * 2) # Space for block and movement
color = rng.randint(1, 9)
question = gen_field(size)
block = [color] * block_size
question = write_block(pos, block, question)
answer = write_block(pos + block_size, block, gen_field(size))
return {"input": question, "output": answer}
def task_change_to_five(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where all non-zero colors change to 5."""
density = 0.5
question = [rng.randint(1, 9) if rng.random() < density else 0 for _ in range(size)]
answer = [5 if x != 0 else 0 for x in question]
return {"input": question, "output": answer}
def task_recolor_blocks_from_palette(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where blocks are recolored using a color palette."""
# Generate blocks of same size
block_size = rng.randint(2, 4)
blocks = []
pos = 0
while pos + block_size <= size:
if rng.random() < 0.4:
blocks.append(pos)
pos += block_size + 1
else:
pos += 1
# Ensure we have space for palette
while blocks and blocks[-1] + block_size + len(blocks) + 1 >= size:
blocks.pop()
if not blocks:
return None
# Shift blocks right to make room for palette
palette_size = len(blocks)
blocks = [pos + palette_size + 1 for pos in blocks]
# Generate color palette
colors = []
for _ in range(len(blocks)):
while True:
color = rng.randint(1, 9)
if color not in colors:
colors.append(color)
break
# Create question with color palette and blocks
question = gen_field(size)
# Place color palette at start
for i, color in enumerate(colors):
question[i] = color
# Place blocks of color 5
for block_pos in blocks:
for i in range(block_size):
question[block_pos + i] = 5
# Create answer with recolored blocks
answer = question.copy()
for block_idx, block_pos in enumerate(blocks):
color = colors[block_idx]
for i in range(block_size):
answer[block_pos + i] = color
return {"input": question, "output": answer}
def task_duplicate_block_from_seeds(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a block is duplicated from seed pixels."""
block_size = rng.randint(2, 4)
if block_size + 1 >= size:
return None
if size <= 3 + block_size:
return None
# Position block with space for seeds
block_pos = rng.randint(2, size - block_size - 1)
# Decide seed placement
left_seed = rng.random() < 0.5
right_seed = rng.random() < 0.5
if not (left_seed or right_seed):
return None
# Create input
question = gen_field(size)
# Place main block
for i in range(block_size):
question[block_pos + i] = 1
# Place seeds with gaps
seeds = []
if left_seed:
color = rng.randint(1, 9)
question[block_pos - 2] = color
seeds.append(("left", block_pos - 2, color))
if right_seed:
color = rng.randint(1, 9)
question[block_pos + block_size + 1] = color
seeds.append(("right", block_pos + block_size + 1, color))
# Create answer with duplicated blocks
answer = question.copy()
for side, seed_pos, color in seeds:
if side == "left":
# For left seed, blocks end at seed
end_pos = seed_pos
while end_pos >= 0:
start_pos = end_pos - block_size + 1
for pos in range(max(0, start_pos), end_pos + 1):
answer[pos] = color
if start_pos < 1:
break
end_pos = start_pos - 2 # -1 for gap
else: # side == "right"
# For right seed, blocks start at seed
start_pos = seed_pos
while start_pos < size:
for offset in range(min(block_size, size - start_pos)):
answer[start_pos + offset] = color
if start_pos + block_size + 1 >= size:
break
start_pos = start_pos + block_size + 1 # +1 for gap
return {"input": question, "output": answer}
def task_fill_from_pixel(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a pixel fills in one direction until hitting another pixel."""
block_size = rng.randint(3, 6)
if block_size >= size - 2:
return None
# Position block with space for seed
block_pos = rng.randint(1, size - block_size - 1)
# Create input
question = gen_field(size)
# Place main block
block_color = rng.randint(1, 9)
for i in range(block_size):
question[block_pos + i] = block_color
# Place seed pixel and determine fill direction
seed_color = rng.randint(1, 9)
while seed_color == block_color:
seed_color = rng.randint(1, 9)
is_left = rng.random() < 0.5
if is_left:
question[block_pos - 1] = seed_color
else:
question[block_pos + block_size] = seed_color
# Create answer with fill
answer = question.copy()
if is_left:
# Fill from seed to left border
for i in range(block_pos):
answer[i] = seed_color
else:
# Fill from seed to right border
for i in range(block_pos + block_size, size):
answer[i] = seed_color
return {"input": question, "output": answer}
def task_mark_size_two_blocks(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where size-2 blocks are marked with surrounding pixels."""
blocks = []
pos = 0
# Generate blocks with minimum gap of 2
while pos < size:
if rng.random() < 0.4:
block_size = rng.randint(1, 3)
# Check if we have space for block and potential markers
needed_space = block_size + (2 if block_size == 2 else 0)
if pos + needed_space < size:
blocks.append((pos, block_size))
pos += block_size + 2 # Minimum gap of 2
pos += 1
if len(blocks) < 2:
return None
# Verify gaps between blocks (including markers)
valid = True
for i in range(len(blocks) - 1):
pos1, size1 = blocks[i]
pos2, _ = blocks[i + 1]
needed_gap = 3 if size1 == 2 else 2
if pos2 - (pos1 + size1) < needed_gap:
valid = False
break
if not valid:
return None
# Create input with blocks
question = gen_field(size)
for pos, block_size in blocks:
# Place block
for i in range(block_size):
question[pos + i] = 1
# Create answer with markers
answer = question.copy()
for pos, block_size in blocks:
if block_size == 2:
# Add markers for size 2 blocks
if pos > 0:
answer[pos - 1] = 3
if pos + block_size < size:
answer[pos + block_size] = 3
return {"input": question, "output": answer}
def task_fill_until_collision(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where pixels fill empty space until collision."""
# At least 4 positions for meaningful puzzle
if size < 4:
return None
is_left = rng.random() < 0.5
question = gen_field(size)
# Place the side marker
if is_left:
question[0] = 5
else:
question[size - 1] = 5
# Place 2-4 random pixels
num_pixels = rng.randint(2, 4)
positions = []
if is_left:
# Skip first position
for _ in range(num_pixels):
while True:
pos = rng.randint(1, size - 1)
if pos not in positions:
positions.append(pos)
break
else:
# Skip last position
for _ in range(num_pixels):
while True:
pos = rng.randint(0, size - 2)
if pos not in positions:
positions.append(pos)
break
# Color random pixels
for pos in positions:
question[pos] = rng.randint(1, 9)
positions.sort()
# Create answer
answer = question.copy()
if is_left:
# Fill right from each pixel
prev_pos = 0 # Start from marker
for pos in positions:
color = question[pos]
# Fill from previous position to current
for i in range(prev_pos + 1, pos):
answer[i] = color
prev_pos = pos
else:
# Fill left from each pixel
prev_pos = size - 1 # Start from marker
for pos in reversed(positions):
color = question[pos]
# Fill from current position to previous
for i in range(pos + 1, prev_pos):
answer[i] = color
prev_pos = pos
return {"input": question, "output": answer}
def task_repeat_pattern_full(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where a pattern is repeated to fill the space."""
# Generate initial pattern
pattern_size = rng.randint(2, 5)
pattern = [rng.randint(1, 9) for _ in range(pattern_size)]
# Calculate total size needed for 2 repetitions
double_size = pattern_size * 2
if double_size >= size:
return None
# Create input with 2 repetitions
question = gen_field(size)
for i in range(pattern_size):
question[i] = pattern[i]
question[i + pattern_size] = pattern[i]
# Create answer with maximum repetitions
answer = gen_field(size)
pos = 0
while pos + pattern_size <= size:
for i in range(pattern_size):
answer[pos + i] = pattern[i]
pos += pattern_size
# Fill remaining space (if any) with pattern elements
for i in range(pos, size):
answer[i] = pattern[i - pos]
return {"input": question, "output": answer}
def task_gravity_weighted_colors(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where color 2 is heavier than color 1 in gravity."""
# Generate random field with only colors 1 and 2
question = [rng.randint(1, 2) if rng.random() < 0.5 else 0 for _ in range(size)]
# Count colors
count_1 = sum(1 for x in question if x == 1)
count_2 = sum(1 for x in question if x == 2)
# Create answer with sorted colors
answer = gen_field(size)
# Place heavier color 2 first
for i in range(count_2):
answer[i] = 2
# Then place color 1
for i in range(count_1):
answer[count_2 + i] = 1
return {"input": question, "output": answer}
def task_color_left_half_blocks(rng: Random, size: int) -> Optional[dict[str, list[int]]]:
"""Generate a task where left half of blocks are colored differently."""
pos = 0
question = gen_field(size)
blocks = []
# Generate blocks with gap 1
while pos < size:
if rng.random() < 0.4:
block_size = rng.randint(2, 8)
if pos + block_size >= size:
break
blocks.append((pos, block_size))
for i in range(block_size):
question[pos + i] = 2
pos += block_size + 1 # block size + gap
else:
pos += 1
if len(blocks) < 2:
return None
# Create answer with half-colored blocks
answer = question.copy()
for pos, block_size in blocks:
half_size = block_size // 2
for i in range(half_size):
answer[pos + i] = 8
return {"input": question, "output": answer}
def task_mirror(task_result: Optional[dict[str, list[int]]]) -> Optional[dict[str, list[int]]]:
"""Mirror the input and output arrays of a task result."""
if task_result is None:
return None
return {"input": list(reversed(task_result["input"])), "output": list(reversed(task_result["output"]))}
def task_inverse(task_result: Optional[dict[str, list[int]]]) -> Optional[dict[str, list[int]]]:
"""Swap the input and output arrays of a task result."""
if task_result is None:
return None
return {"input": task_result["output"], "output": task_result["input"]}
def task_identity(task_result: Optional[dict[str, list[int]]]) -> Optional[dict[str, list[int]]]:
"""Return the task result unchanged."""
return task_result
# Table of all ARC 1D task functions with their parameters
ARC_1D_TASKS = {
# Move tasks - right direction
"move_1pix_solid_right": (task_move_n_pix, {"move_pix": 1, "solid": True}),
"move_2pix_solid_right": (task_move_n_pix, {"move_pix": 2, "solid": True}),
"move_3pix_solid_right": (task_move_n_pix, {"move_pix": 3, "solid": True}),
"move_4pix_solid_right": (task_move_n_pix, {"move_pix": 4, "solid": True}),
"move_1pix_colorful_right": (task_move_n_pix, {"move_pix": 1, "solid": False}),
"move_2pix_colorful_right": (task_move_n_pix, {"move_pix": 2, "solid": False}),
"move_3pix_colorful_right": (task_move_n_pix, {"move_pix": 3, "solid": False}),
"move_4pix_colorful_right": (task_move_n_pix, {"move_pix": 4, "solid": False}),
# Move tasks - left direction (mirrored)
"move_1pix_solid_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix(rng, size, **kwargs)),
{"move_pix": 1, "solid": True},
),
"move_2pix_solid_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix(rng, size, **kwargs)),
{"move_pix": 2, "solid": True},
),
"move_3pix_solid_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix(rng, size, **kwargs)),
{"move_pix": 3, "solid": True},
),
"move_4pix_solid_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix(rng, size, **kwargs)),
{"move_pix": 4, "solid": True},
),
"move_1pix_colorful_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix(rng, size, **kwargs)),
{"move_pix": 1, "solid": False},
),
"move_2pix_colorful_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix(rng, size, **kwargs)),
{"move_pix": 2, "solid": False},
),
"move_3pix_colorful_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix(rng, size, **kwargs)),
{"move_pix": 3, "solid": False},
),
"move_4pix_colorful_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix(rng, size, **kwargs)),
{"move_pix": 4, "solid": False},
),
# Move wrapped tasks - right direction
"move_1pix_solid_wrapped_right": (task_move_n_pix_wrapped, {"move_pix": 1, "solid": True}),
"move_2pix_solid_wrapped_right": (task_move_n_pix_wrapped, {"move_pix": 2, "solid": True}),
"move_3pix_solid_wrapped_right": (task_move_n_pix_wrapped, {"move_pix": 3, "solid": True}),
"move_4pix_solid_wrapped_right": (task_move_n_pix_wrapped, {"move_pix": 4, "solid": True}),
"move_1pix_colorful_wrapped_right": (task_move_n_pix_wrapped, {"move_pix": 1, "solid": False}),
"move_2pix_colorful_wrapped_right": (task_move_n_pix_wrapped, {"move_pix": 2, "solid": False}),
"move_3pix_colorful_wrapped_right": (task_move_n_pix_wrapped, {"move_pix": 3, "solid": False}),
"move_4pix_colorful_wrapped_right": (task_move_n_pix_wrapped, {"move_pix": 4, "solid": False}),
# Move wrapped tasks - left direction (mirrored)
"move_1pix_solid_wrapped_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix_wrapped(rng, size, **kwargs)),
{"move_pix": 1, "solid": True},
),
"move_2pix_solid_wrapped_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix_wrapped(rng, size, **kwargs)),
{"move_pix": 2, "solid": True},
),
"move_3pix_solid_wrapped_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix_wrapped(rng, size, **kwargs)),
{"move_pix": 3, "solid": True},
),
"move_4pix_solid_wrapped_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix_wrapped(rng, size, **kwargs)),
{"move_pix": 4, "solid": True},
),
"move_1pix_colorful_wrapped_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix_wrapped(rng, size, **kwargs)),
{"move_pix": 1, "solid": False},
),
"move_2pix_colorful_wrapped_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix_wrapped(rng, size, **kwargs)),
{"move_pix": 2, "solid": False},
),
"move_3pix_colorful_wrapped_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix_wrapped(rng, size, **kwargs)),
{"move_pix": 3, "solid": False},
),
"move_4pix_colorful_wrapped_left": (
lambda rng, size, **kwargs: task_mirror(task_move_n_pix_wrapped(rng, size, **kwargs)),
{"move_pix": 4, "solid": False},
),
# Gravity tasks - right direction
"gravity_right": (task_gravity, {}),
"gravity_counting_right": (task_gravity_counting, {}),
"gravity_antigravity_right": (task_gravity_antigravity, {}),
"gravity_one_step_right": (task_gravity_one_step, {}),
"gravity_weighted_colors_right": (task_gravity_weighted_colors, {}),
# Gravity tasks - left direction (mirrored)
"gravity_left": (lambda rng, size, **kwargs: task_mirror(task_gravity(rng, size, **kwargs)), {}),
"gravity_counting_left": (lambda rng, size, **kwargs: task_mirror(task_gravity_counting(rng, size, **kwargs)), {}),
"gravity_antigravity_left": (
lambda rng, size, **kwargs: task_mirror(task_gravity_antigravity(rng, size, **kwargs)),
{},
),
"gravity_one_step_left": (lambda rng, size, **kwargs: task_mirror(task_gravity_one_step(rng, size, **kwargs)), {}),
"gravity_weighted_colors_left": (
lambda rng, size, **kwargs: task_mirror(task_gravity_weighted_colors(rng, size, **kwargs)),
{},
),
# Block tasks
"block_touch_dot": (task_block_touch_dot, {}),
"block_touch_dot_1pix": (task_block_touch_dot_n_pix, {"move_pix": 1}),
"block_touch_dot_2pix": (task_block_touch_dot_n_pix, {"move_pix": 2}),
"block_touch_dot_3pix": (task_block_touch_dot_n_pix, {"move_pix": 3}),
"block_touch_dot_4pix": (task_block_touch_dot_n_pix, {"move_pix": 4}),
"block_scale_to_dot": (task_block_scale_to_dot, {}),
"block_and_noise_remove": (task_block_and_noise_remove, {}),
"block_and_noise_remove_inside": (task_block_and_noise_remove_inside, {}),
"move_block_by_own_size": (task_move_block_by_own_size, {}),
# Pattern tasks
"two_points_and_fill": (task_two_points_and_fill, {}),
"two_points_and_fill_inv": (
lambda rng, size, **kwargs: task_inverse(task_two_points_and_fill(rng, size, **kwargs)),
{},
),
"copy_block_to_dots": (task_copy_block_to_dots, {}),
"copy_block_to_dots_colors": (task_copy_block_to_dots_colors, {}),
"repeat_pattern_full": (task_repeat_pattern_full, {}),
# Reflection tasks
"reflect_block_with_border_pixel": (task_reflect_block_with_border_pixel, {}),
"reflect_block_random": (task_reflect_block_with_border_pixel_random, {}),
"reflect_block_around_dot": (task_reflect_block_around_dot, {}),
# Color tasks
"paint_biggest_block": (task_paint_biggest_block, {}),
"recolor_blocks_by_size": (task_recolor_blocks_by_size, {}),
"change_to_five": (task_change_to_five, {}),
"recolor_blocks_from_palette": (task_recolor_blocks_from_palette, {}),
"color_left_half_blocks": (task_color_left_half_blocks, {}),
# Sorting tasks
"sort_blocks_by_size": (task_sort_blocks_by_size, {}),
"sort_complete_sequence": (task_sort_complete_sequence, {}),
# Fill tasks
"duplicate_block_from_seeds": (task_duplicate_block_from_seeds, {}),
"fill_from_pixel": (task_fill_from_pixel, {}),
"fill_until_collision": (task_fill_until_collision, {}),
# Marking tasks
"mark_size_two_blocks": (task_mark_size_two_blocks, {}),
}