File size: 16,171 Bytes
9485e3f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
#!/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)