Problem 542942 · medium · Phase 05 Advanced Algorithms & Graphs

Picket Colouring

dynamic programming · 1-D dp · counting · state machine · modular arithmetic

A fence has n pickets in a row and you own paint in k colours. Any picket may be any colour, except that three consecutive pickets must never share a colour. Return the number of valid colourings modulo 10**9 + 7.

Examples

Input:  n = 3, k = 2
Output: 6
Explanation: of the 8 two-colour patterns only AAA and BBB are forbidden.

Input:  n = 1, k = 5
Output: 5

Constraints

  • 1 <= n <= 10**5, 1 <= k <= 10**5
  • Target complexity: O(n) time.

Goals

  • Split the count into 'same as previous' and 'different from previous' states
  • Derive transitions from a no-three-in-a-row rule
Starting Python…