#!/usr/bin/env python3 """Generate cold-start reasoning data (Stage 2b) — 5K problem→reasoning→solution→test pairs.""" import json import random import os from pathlib import Path SEED = 42 random.seed(SEED) PROBLEMS = [ { "task": "two_sum", "instruction": "Write a function that finds two numbers in a list that add up to a target sum.", "reasoning": "Let me solve this step by step.\n1. I need to find two numbers in the array that sum to target.\n2. A brute force approach would check every pair O(n²), but I can do better.\n3. I'll use a hash map to store numbers I've seen and check if target - current exists.\n4. This gives O(n) time and O(n) space, which is optimal.", "solution": "def two_sum(nums, target):\n seen = {}\n for i, n in enumerate(nums):\n complement = target - n\n if complement in seen:\n return [seen[complement], i]\n seen[n] = i\n return []", "tests": "assert two_sum([2,7,11,15], 9) == [0,1]\nassert two_sum([3,2,4], 6) == [1,2]\nassert two_sum([3,3], 6) == [0,1]\nassert two_sum([1,2,3], 7) == []" }, { "task": "reverse_string", "instruction": "Write a function that reverses a string in-place.", "reasoning": "Let me think about reversing a string.\n1. To reverse in-place, I need two pointers: one at start, one at end.\n2. I swap characters at both pointers, then move them inward.\n3. I continue until the pointers meet or cross.\n4. For Python strings (immutable), I'll convert to a list first.", "solution": "def reverse_string(s):\n chars = list(s)\n left, right = 0, len(chars) - 1\n while left < right:\n chars[left], chars[right] = chars[right], chars[left]\n left += 1\n right -= 1\n return ''.join(chars)", "tests": "assert reverse_string('hello') == 'olleh'\nassert reverse_string('') == ''\nassert reverse_string('a') == 'a'\nassert reverse_string('ab') == 'ba'" }, { "task": "is_palindrome", "instruction": "Write a function that checks if a string is a palindrome.", "reasoning": "A palindrome reads the same forwards and backwards.\n1. I'll use two pointers: one at start, one at end.\n2. While left < right, I compare characters and move inward.\n3. If any mismatch, it's not a palindrome.\n4. If all match, it is a palindrome.\n5. I'll skip non-alphanumeric characters for a more robust solution.", "solution": "def is_palindrome(s):\n left, right = 0, len(s) - 1\n while left < right:\n while left < right and not s[left].isalnum():\n left += 1\n while left < right and not s[right].isalnum():\n right -= 1\n if s[left].lower() != s[right].lower():\n return False\n left += 1\n right -= 1\n return True", "tests": "assert is_palindrome('racecar') == True\nassert is_palindrome('hello') == False\nassert is_palindrome('A man, a plan, a canal: Panama') == True\nassert is_palindrome('') == True" }, { "task": "max_subarray", "instruction": "Write a function that finds the maximum sum of any contiguous subarray (Kadane's algorithm).", "reasoning": "Kadane's algorithm is the standard O(n) solution.\n1. I keep track of the current sum and the best sum seen so far.\n2. At each element, I decide: extend the current subarray or start fresh.\n3. This is: current = max(num, current + num).\n4. I update best = max(best, current).\n5. This works because a subarray ending at position i is either the element itself or the element plus the best subarray ending at i-1.", "solution": "def max_subarray(nums):\n best = nums[0]\n current = nums[0]\n for n in nums[1:]:\n current = max(n, current + n)\n best = max(best, current)\n return best", "tests": "assert max_subarray([-2,1,-3,4,-1,2,1,-5,4]) == 6\nassert max_subarray([1]) == 1\nassert max_subarray([5,4,-1,7,8]) == 23\nassert max_subarray([-1,-2,-3]) == -1" }, { "task": "merge_intervals", "instruction": "Write a function that merges overlapping intervals.", "reasoning": "Let me solve interval merging.\n1. First, sort intervals by start time.\n2. Then iterate through sorted intervals.\n3. If the current interval overlaps with the last merged interval, merge them.\n4. Overlap condition: current.start <= last_merged.end.\n5. Merged end = max(last_merged.end, current.end).\n6. If no overlap, add current to result.", "solution": "def merge_intervals(intervals):\n if not intervals:\n return []\n intervals.sort(key=lambda x: x[0])\n merged = [intervals[0]]\n for start, end in intervals[1:]:\n if start <= merged[-1][1]:\n merged[-1] = (merged[-1][0], max(merged[-1][1], end))\n else:\n merged.append([start, end])\n return merged", "tests": "assert merge_intervals([[1,3],[2,6],[8,10],[15,18]]) == [[1,6],[8,10],[15,18]]\nassert merge_intervals([[1,4],[4,5]]) == [[1,5]]\nassert merge_intervals([[1,2]]) == [[1,2]]\nassert merge_intervals([]) == []" }, { "task": "valid_parentheses", "instruction": "Write a function that checks if a string of brackets is valid.", "reasoning": "I'll use a stack for this classic problem.\n1. Map closing brackets to their corresponding opening brackets.\n2. Iterate through the string.\n3. If it's an opening bracket, push onto stack.\n4. If it's a closing bracket, check if the top of stack matches.\n5. If not, it's invalid.\n6. At the end, the stack should be empty.", "solution": "def valid_parentheses(s):\n pairs = {')': '(', '}': '{', ']': '['}\n stack = []\n for ch in s:\n if ch in pairs:\n if not stack or stack[-1] != pairs[ch]:\n return False\n stack.pop()\n else:\n stack.append(ch)\n return len(stack) == 0", "tests": "assert valid_parentheses('()') == True\nassert valid_parentheses('()[]{}') == True\nassert valid_parentheses('(]') == False\nassert valid_parentheses('([)]') == False\nassert valid_parentheses('{[]}') == True" }, { "task": "lru_cache", "instruction": "Design an LRU (Least Recently Used) cache.", "reasoning": "An LRU cache needs O(1) get and put operations.\n1. I'll use a dict for key-value lookups and a doubly-linked list for order.\n2. The dict maps keys to nodes in the linked list.\n3. On get: move accessed node to head (most recent).\n4. On put: add new node to head, evict tail if over capacity.\n5. Using a sentinel head and tail simplifies edge cases.", "solution": "class LRUCache:\n def __init__(self, capacity):\n self.cap = capacity\n self.cache = {}\n self.head = {'key': None, 'val': None, 'prev': None, 'next': None}\n self.tail = {'key': None, 'val': None, 'prev': self.head, 'next': None}\n self.head['next'] = self.tail\n\n def _remove(self, node):\n node['prev']['next'] = node['next']\n node['next']['prev'] = node['prev']\n\n def _add_to_head(self, node):\n node['next'] = self.head['next']\n node['prev'] = self.head\n self.head['next']['prev'] = node\n self.head['next'] = node\n\n def get(self, key):\n if key not in self.cache:\n return -1\n node = self.cache[key]\n self._remove(node)\n self._add_to_head(node)\n return node['val']\n\n def put(self, key, value):\n if key in self.cache:\n node = self.cache[key]\n node['val'] = value\n self._remove(node)\n self._add_to_head(node)\n else:\n node = {'key': key, 'val': value, 'prev': None, 'next': None}\n self.cache[key] = node\n self._add_to_head(node)\n if len(self.cache) > self.cap:\n lru = self.tail['prev']\n self._remove(lru)\n del self.cache[lru['key']]", "tests": "cache = LRUCache(2)\ncache.put(1,1)\ncache.put(2,2)\nassert cache.get(1) == 1\ncache.put(3,3)\nassert cache.get(2) == -1\ncache.put(4,4)\nassert cache.get(1) == -1\nassert cache.get(3) == 3\nassert cache.get(4) == 4" }, { "task": "word_count", "instruction": "Write a function that counts word frequencies in a text.", "reasoning": "I'll split the text into words, normalize them, and count.\n1. Split by whitespace and remove punctuation.\n2. Convert to lowercase for case-insensitive counting.\n3. Use a dict to store word->count.\n4. Return the dict sorted by frequency descending.\n5. Handle empty input gracefully.", "solution": "def word_count(text):\n import re\n words = re.findall(r'\\w+', text.lower())\n freq = {}\n for w in words:\n freq[w] = freq.get(w, 0) + 1\n return dict(sorted(freq.items(), key=lambda x: -x[1]))", "tests": "assert word_count('hello world hello') == {'hello': 2, 'world': 1}\nassert word_count('') == {}\nassert word_count('Go go go') == {'go': 3}" }, { "task": "binary_search", "instruction": "Write a function that performs binary search on a sorted list.", "reasoning": "Binary search is O(log n) on sorted arrays.\n1. Set left = 0, right = len(arr) - 1.\n2. While left <= right, check the middle element.\n3. If middle == target, return its index.\n4. If middle < target, search right half.\n5. If middle > target, search left half.\n6. If not found, return -1.", "solution": "def binary_search(arr, target):\n left, right = 0, len(arr) - 1\n while left <= right:\n mid = (left + right) // 2\n if arr[mid] == target:\n return mid\n elif arr[mid] < target:\n left = mid + 1\n else:\n right = mid - 1\n return -1", "tests": "assert binary_search([1,2,3,4,5], 3) == 2\nassert binary_search([1,2,3,4,5], 6) == -1\nassert binary_search([], 1) == -1\nassert binary_search([1], 1) == 0" }, { "task": "linked_list_reverse", "instruction": "Write a function to reverse a linked list.", "reasoning": "I'll reverse the linked list iteratively.\n1. Use three pointers: prev, current, next.\n2. For each node, save next, point current.next to prev, then advance.\n3. At the end, prev will be the new head.\n4. This is O(n) time and O(1) space.", "solution": "def reverse_linked_list(head):\n prev = None\n current = head\n while current:\n next_node = current['next']\n current['next'] = prev\n prev = current\n current = next_node\n return prev", "tests": "# Using dict as node for testing\nn3 = {'val': 3, 'next': None}\nn2 = {'val': 2, 'next': n3}\nn1 = {'val': 1, 'next': n2}\nreversed_head = reverse_linked_list(n1)\nassert reversed_head['val'] == 3\nassert reversed_head['next']['val'] == 2\nassert reversed_head['next']['next']['val'] == 1" }, { "task": "anagram_check", "instruction": "Write a function that checks if two strings are anagrams.", "reasoning": "Two strings are anagrams if they have the same characters with the same frequencies.\n1. First approach: sort both strings and compare O(n log n).\n2. Better approach: count character frequencies with a dict O(n).\n3. If the frequency dicts are equal, they're anagrams.\n4. Return False early if lengths differ.", "solution": "def is_anagram(s, t):\n if len(s) != len(t):\n return False\n counts = {}\n for ch in s:\n counts[ch] = counts.get(ch, 0) + 1\n for ch in t:\n if ch not in counts:\n return False\n counts[ch] -= 1\n if counts[ch] < 0:\n return False\n return True", "tests": "assert is_anagram('listen', 'silent') == True\nassert is_anagram('hello', 'world') == False\nassert is_anagram('', '') == True\nassert is_anagram('aabb', 'bbaa') == True" }, { "task": "find_missing", "instruction": "Write a function that finds the missing number in an array from 0 to n.", "reasoning": "I can use the mathematical sum formula.\n1. The expected sum from 0 to n is n*(n+1)//2.\n2. Sum the actual numbers in the array.\n3. The difference is the missing number.\n4. This is O(n) time and O(1) space, no extra data structures needed.", "solution": "def find_missing(nums):\n n = len(nums)\n expected = n * (n + 1) // 2\n actual = sum(nums)\n return expected - actual", "tests": "assert find_missing([3,0,1]) == 2\nassert find_missing([0,1]) == 2\nassert find_missing([9,6,4,2,3,5,7,0,1]) == 8\nassert find_missing([0]) == 1" }, ] DEBUG_PROBLEMS = [ { "task": "debug_fizzbuzz", "instruction": "Find and fix the bug in this FizzBuzz implementation: for i in range(1, 16):\n if i % 3 == 0:\n print('Fizz')\n elif i % 5 == 0:\n print('Buzz')\n else:\n print(i)", "reasoning": "Let me analyze the bug step by step.\n1. The problem: numbers divisible by both 3 and 5 should print 'FizzBuzz'.\n2. The current code uses elif, so when i % 3 == 0 is True, it prints 'Fizz' and skips the i % 5 check.\n3. This means i=15 prints 'Fizz' instead of 'FizzBuzz'.\n4. Fix: check the combined condition first (i % 15 == 0 or i % 3 == 0 and i % 5 == 0).", "solution": "def fizzbuzz(n):\n for i in range(1, n+1):\n if i % 15 == 0:\n print('FizzBuzz')\n elif i % 3 == 0:\n print('Fizz')\n elif i % 5 == 0:\n print('Buzz')\n else:\n print(i)", "tests": "# Expected output for n=15\n# 1,2,Fizz,4,Buzz,Fizz,7,8,Fizz,Buzz,11,Fizz,13,14,FizzBuzz" }, { "task": "debug_mutable_default", "instruction": "Find the bug: def add_item(item, items=[]):\n items.append(item)\n return items", "reasoning": "This is a classic Python mutable default argument bug.\n1. Default arguments are evaluated once at function definition time.\n2. All calls that use the default will share the same list.\n3. So add_item(1) returns [1], add_item(2) returns [1,2], not [2].\n4. Fix: use None as default and create a new list inside.", "solution": "def add_item(item, items=None):\n if items is None:\n items = []\n items.append(item)\n return items", "tests": "assert add_item(1) == [1]\nassert add_item(2) == [2]\nassert add_item(3, [0]) == [0,3]" }, ] def generate_coldstart(output_path, num_examples=5000): os.makedirs(os.path.dirname(output_path), exist_ok=True) examples = [] all_problems = PROBLEMS + DEBUG_PROBLEMS while len(examples) < num_examples: p = random.choice(all_problems) examples.append({ "input": p["instruction"], "reasoning": p["reasoning"], "output": p["solution"], "tests": p["tests"], "source": "cold_start", "task": p["task"], }) # Trim to exact count examples = examples[:num_examples] random.shuffle(examples) with open(output_path, 'w') as f: for ex in examples: f.write(json.dumps(ex) + '\n') print(f"Generated {len(examples)} cold-start examples -> {output_path}") return examples def save_as_jsonl(examples, output_path): """Save examples as JSONL for the ColdStartDataset.""" with open(output_path, 'w') as f: for ex in examples: f.write(json.dumps(ex) + '\n') stats = {"total": len(examples), "tasks": {}} for ex in examples: t = ex.get("task", "unknown") stats["tasks"][t] = stats["tasks"].get(t, 0) + 1 print(f"Saved {len(examples)} examples to {output_path}") print(f"Task distribution: {stats['tasks']}") if __name__ == '__main__': import argparse parser = argparse.ArgumentParser() parser.add_argument('--output', type=str, default='/FSI_Edge/data/cold_start.jsonl') parser.add_argument('--num', type=int, default=5000) args = parser.parse_args() generate_coldstart(args.output, args.num)