AdventOfCode2025/day03.py at main · willkill07/AdventOfCode2025 · GitHub

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
import numpy as np
import nvtx
import warp as wp

wp.init()


@wp.kernel
def mark(a: wp.array(dtype=wp.int8), o: wp.array(dtype=wp.int8), W: int, K: int):
    r = wp.tid()
    W = wp.int32(W)
    K = wp.int32(K)
    s = r * W
    b = r * K
    d = wp.int32(W - K)
    k = wp.int32(0)
    for i in range(W):
        v = a[s + i]
        while d > 0 and k > 0:
            if o[b + k - 1] >= v:
                break
            k -= 1
            d -= 1
        if k < K:
            o[b + k] = v
            k += 1
        else:
            d -= 1


def run(d: wp.array(dtype=wp.int8), N: int, W: int, K: int) -> int:
    o = wp.zeros(N * K, dtype=wp.int8)
    wp.launch(
        mark,
        dim=N,
        inputs=[d, o, W, K],
    )
    return (
        o.numpy().reshape(N, K).astype(np.int64)
        @ (10 ** np.arange(K - 1, -1, -1, dtype=np.int64))
    ).sum()


@nvtx.annotate("Input Parsing")
def parse() -> tuple[wp.array, int, int]:
    with open("inputs/day03.in", "rb") as f:
        d = f.read()
    W = d.index(b"\n")
    d = d.replace(b"\n", b"")
    return (
        wp.array((np.frombuffer(d, dtype=np.uint8) - 48).astype(np.int8)),
        len(d) // W,
        W,
    )


@nvtx.annotate("Part 1")
def part1(d: wp.array(dtype=wp.int8), N: int, W: int) -> int:
    return run(d, N, W, 2)


@nvtx.annotate("Part 2")
def part2(d: wp.array(dtype=wp.int8), N: int, W: int) -> int:
    return run(d, N, W, 12)


@nvtx.annotate("Day 03")
def main() -> tuple[int, int]:
    d = parse()
    return part1(*d), part2(*d)


if __name__ == "__main__":
    print(*main())