when not declared ATCODER_YULY3HEADER_HPP: const ATCODER_YULY3HEADER_HPP* = 1 import algorithm, bitops, deques, heapqueue, math, macros, sets, sequtils, strformat, strutils, sugar, tables proc transLastStmt(n, res, bracketExpr: NimNode): (NimNode, NimNode, NimNode) = # Looks for the last statement of the last statement, etc... case n.kind of nnkIfExpr, nnkIfStmt, nnkTryStmt, nnkCaseStmt: result[0] = copyNimTree(n) result[1] = copyNimTree(n) result[2] = copyNimTree(n) for i in ord(n.kind == nnkCaseStmt)..= 1: (result[0][^1], result[1][^1], result[2][^1]) = transLastStmt(n[^1], res, bracketExpr) of nnkTableConstr: result[1] = n[0][0] result[2] = n[0][1] if bracketExpr.len == 1: bracketExpr.add([newCall(bindSym"typeof", newEmptyNode()), newCall( bindSym"typeof", newEmptyNode())]) template adder(res, k, v) = res[k] = v result[0] = getAst(adder(res, n[0][0], n[0][1])) of nnkCurly: result[2] = n[0] if bracketExpr.len == 1: bracketExpr.add(newCall(bindSym"typeof", newEmptyNode())) template adder(res, v) = res.incl(v) result[0] = getAst(adder(res, n[0])) else: result[2] = n if bracketExpr.len == 1: bracketExpr.add(newCall(bindSym"typeof", newEmptyNode())) template adder(res, v) = res.add(v) result[0] = getAst(adder(res, n)) macro collect*(init, body: untyped): untyped = runnableExamples: import sets, tables let data = @["bird", "word"] ## seq: let k = collect(newSeq): for i, d in data.pairs: if i mod 2 == 0: d assert k == @["bird"] ## seq with initialSize: let x = collect(newSeqOfCap(4)): for i, d in data.pairs: if i mod 2 == 0: d assert x == @["bird"] ## HashSet: let y = initHashSet.collect: for d in data.items: {d} assert y == data.toHashSet ## Table: let z = collect(initTable(2)): for i, d in data.pairs: {i: d} assert z == {0: "bird", 1: "word"}.toTable let res = genSym(nskVar, "collectResult") expectKind init, {nnkCall, nnkIdent, nnkSym} let bracketExpr = newTree(nnkBracketExpr, if init.kind == nnkCall: init[0] else: init) let (resBody, keyType, valueType) = transLastStmt(body, res, bracketExpr) if bracketExpr.len == 3: bracketExpr[1][1] = keyType bracketExpr[2][1] = valueType else: bracketExpr[1][1] = valueType let call = newTree(nnkCall, bracketExpr) if init.kind == nnkCall: for i in 1 ..< init.len: call.add init[i] result = newTree(nnkStmtListExpr, newVarStmt(res, call), resBody, res) proc nextString*(f: auto = stdin): string = var get = false result = "" while true: let c = f.readChar if c.int > ' '.int: get = true result.add(c) elif get: return proc nextInt*(f: auto = stdin): int = parseInt(f.nextString) proc nextFloat*(f: auto = stdin): float = parseFloat(f.nextString) proc chmax*[T](n: var T, m: T) {.inline.} = n = max(n, m) proc chmin*[T](n: var T, m: T) {.inline.} = n = min(n, m) proc `%=`*[T: SomeInteger](n: var T, m: T) {.inline.} = n = floorMod(n, m) proc `|=`*[T: SomeInteger or bool](n: var T, m: T) {.inline.} = n = n or m proc `&=`*[T: SomeInteger or bool](n: var T, m: T) {.inline.} = n = n and m proc `^=`*[T: SomeInteger or bool](n: var T, m: T) {.inline.} = n = n xor m proc `<<=`*[T: SomeInteger](n: var T, m: T) {.inline.} = n = n shl m proc `>>=`*[T: SomeInteger](n: var T, m: T) {.inline.} = n = n shr m when isMainModule: var N = nextInt() var A: array[1000001, int] for i in 0..