Improve documentation

[smuralee]

Improve the documentation for the code samples by referencing the problem statement in the comments. This will provide context for the reader to understand the solution/ approach.